From 8a4ffea9bef3a560b653088a9f43936b2b4a0fcb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adam=20Chuda=C5=9B?= Date: Wed, 18 Sep 2024 19:57:48 +0200 Subject: [PATCH] In-memory shard UId mapping --- chain/chain/src/flat_storage_creator.rs | 10 +- chain/chain/src/garbage_collection.rs | 1 - chain/chain/src/runtime/mod.rs | 5 + core/primitives/src/shard_layout.rs | 14 -- core/store/src/cold_storage.rs | 5 +- core/store/src/columns.rs | 10 +- core/store/src/flat/inlining_migration.rs | 20 ++- core/store/src/flat/store_helper.rs | 4 + core/store/src/genesis/initialization.rs | 6 +- core/store/src/lib.rs | 8 +- core/store/src/test_utils.rs | 7 + core/store/src/trie/from_flat.rs | 8 +- core/store/src/trie/mem/parallel_loader.rs | 15 +- core/store/src/trie/mod.rs | 6 +- .../src/trie/prefetching_trie_storage.rs | 26 +-- core/store/src/trie/shard_tries.rs | 75 +++++---- core/store/src/trie/state_parts.rs | 7 +- core/store/src/trie/state_reader.rs | 153 ++++++++++++++++++ core/store/src/trie/trie_storage.rs | 60 ++----- core/store/src/trie/trie_tests.rs | 52 ++++-- .../src/tests/client/state_snapshot.rs | 3 + nearcore/src/entity_debug.rs | 32 ++-- nearcore/src/metrics.rs | 19 +-- .../src/estimator_context.rs | 10 +- runtime/runtime/src/prefetch.rs | 3 + tools/database/src/analyze_contract_sizes.rs | 8 +- tools/database/src/analyze_delayed_receipt.rs | 3 + tools/database/src/state_perf.rs | 7 +- tools/fork-network/src/cli.rs | 4 +- tools/state-viewer/src/commands.rs | 11 +- .../src/trie_iteration_benchmark.rs | 4 +- 31 files changed, 397 insertions(+), 199 deletions(-) create mode 100644 core/store/src/trie/state_reader.rs diff --git a/chain/chain/src/flat_storage_creator.rs b/chain/chain/src/flat_storage_creator.rs index 70064dc908a..af4d13bd3ef 100644 --- a/chain/chain/src/flat_storage_creator.rs +++ b/chain/chain/src/flat_storage_creator.rs @@ -25,6 +25,7 @@ use near_store::flat::{ FlatStorageCreationStatus, FlatStorageManager, FlatStorageReadyStatus, FlatStorageStatus, NUM_PARTS_IN_ONE_STEP, STATE_PART_MEMORY_LIMIT, }; +use near_store::trie::StateReader; use near_store::Store; use near_store::{Trie, TrieDBStorage, TrieTraversalItem}; use std::collections::HashMap; @@ -94,8 +95,9 @@ impl FlatStorageShardCreator { part_id: PartId, progress: Arc, result_sender: Sender, + state_reader: StateReader, ) { - let trie_storage = TrieDBStorage::new(store.clone(), shard_uid); + let trie_storage = TrieDBStorage::new(store.clone(), shard_uid, state_reader); let trie = Trie::new(Arc::new(trie_storage), state_root, None); let path_begin = trie.find_state_part_boundary(part_id.idx, part_id.total).unwrap(); let path_end = trie.find_state_part_boundary(part_id.idx + 1, part_id.total).unwrap(); @@ -195,7 +197,8 @@ impl FlatStorageShardCreator { let store = self.runtime.store().clone(); let epoch_id = self.epoch_manager.get_epoch_id(&block_hash)?; let shard_uid = self.epoch_manager.shard_id_to_uid(shard_id, &epoch_id)?; - let trie_storage = TrieDBStorage::new(store, shard_uid); + let state_reader = self.runtime.get_tries().get_state_reader(); + let trie_storage = TrieDBStorage::new(store, shard_uid, state_reader); let state_root = *chain_store.get_chunk_extra(&block_hash, &shard_uid)?.state_root(); let trie = Trie::new(Arc::new(trie_storage), state_root, None); @@ -226,6 +229,7 @@ impl FlatStorageShardCreator { fetching_state_status, )) => { let store = self.runtime.store().clone(); + let state_reader = self.runtime.get_tries().get_state_reader(); let block_hash = fetching_state_status.block_hash; let start_part_id = fetching_state_status.part_id; let num_parts_in_step = fetching_state_status.num_parts_in_step; @@ -251,6 +255,7 @@ impl FlatStorageShardCreator { let inner_progress = progress.clone(); let inner_sender = self.fetched_parts_sender.clone(); let inner_threads_used = self.metrics.threads_used(); + let state_reader = state_reader.clone(); thread_pool.spawn(move || { inner_threads_used.inc(); Self::fetch_state_part( @@ -260,6 +265,7 @@ impl FlatStorageShardCreator { PartId::new(part_id, num_parts), inner_progress, inner_sender, + state_reader, ); inner_threads_used.dec(); }) diff --git a/chain/chain/src/garbage_collection.rs b/chain/chain/src/garbage_collection.rs index 042835e7884..599119a112a 100644 --- a/chain/chain/src/garbage_collection.rs +++ b/chain/chain/src/garbage_collection.rs @@ -993,7 +993,6 @@ impl<'a> ChainStoreUpdate<'a> { | DBCol::FlatStorageStatus | DBCol::Misc | DBCol::_ReceiptIdToShardId - | DBCol::ShardUIdMapping => unreachable!(), } self.merge(store_update); diff --git a/chain/chain/src/runtime/mod.rs b/chain/chain/src/runtime/mod.rs index e79b8e8ea8e..d03b0a65af3 100644 --- a/chain/chain/src/runtime/mod.rs +++ b/chain/chain/src/runtime/mod.rs @@ -40,6 +40,7 @@ use near_primitives::views::{ use near_store::config::StateSnapshotType; use near_store::flat::FlatStorageManager; use near_store::metadata::DbKind; +use near_store::trie::StateReader; use near_store::{ ApplyStatePartResult, DBCol, ShardTries, StateSnapshotConfig, Store, Trie, TrieConfig, TrieUpdate, WrappedTrieChanges, COLD_HEAD_KEY, @@ -101,10 +102,14 @@ impl NightshadeRuntime { let trie_viewer = TrieViewer::new(trie_viewer_state_size_limit, max_gas_burnt_view); let flat_storage_manager = FlatStorageManager::new(store.clone()); let shard_uids: Vec<_> = genesis_config.shard_layout.shard_uids().collect(); + // TODO(reshardingV3) Recursively calculate resharding parents for `shard_uids`. + let shard_uids_ancestor_tree = shard_uids.iter().map(|id| (*id, None)).collect(); + let state_reader = StateReader::new(&shard_uids, shard_uids_ancestor_tree); let tries = ShardTries::new( store.clone(), trie_config, &shard_uids, + state_reader, flat_storage_manager, state_snapshot_config, ); diff --git a/core/primitives/src/shard_layout.rs b/core/primitives/src/shard_layout.rs index 7042bdc718c..f15ab81b10d 100644 --- a/core/primitives/src/shard_layout.rs +++ b/core/primitives/src/shard_layout.rs @@ -325,20 +325,6 @@ pub struct ShardUId { pub shard_id: u32, } -/// A wrapper for ShardUId used to denote when we expect the mapped version of the value. -/// -/// ReshardingV3 uses mapping strategy for the State column. -/// It means we access trie storage using a mapped shard UID, that can be the same shard UID, -/// resharding parent shard UID, or any ancestor shard UID in the tree of resharding history. -#[derive(Clone, Copy)] -pub struct MappedShardUId(pub ShardUId); - -impl From for MappedShardUId { - fn from(shard_uid: ShardUId) -> Self { - MappedShardUId(shard_uid) - } -} - impl ShardUId { pub fn single_shard() -> Self { Self { version: 0, shard_id: 0 } diff --git a/core/store/src/cold_storage.rs b/core/store/src/cold_storage.rs index 6cf7a9a0d2b..27e066a52c7 100644 --- a/core/store/src/cold_storage.rs +++ b/core/store/src/cold_storage.rs @@ -1,6 +1,6 @@ use crate::columns::DBKeyType; use crate::db::{ColdDB, COLD_HEAD_KEY, HEAD_KEY}; -use crate::{map_shard_uid, metrics, DBCol, DBTransaction, Database, Store, TrieChanges}; +use crate::{metrics, DBCol, DBTransaction, Database, Store, TrieChanges}; use borsh::BorshDeserialize; use near_primitives::block::{Block, BlockHeader, Tip}; @@ -185,7 +185,8 @@ fn copy_state_from_store( hot_store.get_ser::(DBCol::TrieChanges, &key)?; let Some(trie_changes) = trie_changes else { continue }; - let shard_uid_db_prefix = map_shard_uid(&hot_store, shard_uid).0.to_bytes(); + // TODO(reshardingV3) Likely should use mapped shard UId + let shard_uid_db_prefix = shard_uid.to_bytes(); for op in trie_changes.insertions() { let key = join_two_keys(&shard_uid_db_prefix, op.hash().as_bytes()); let value = op.payload().to_vec(); diff --git a/core/store/src/columns.rs b/core/store/src/columns.rs index 208f590ab0b..1ac720a5e2d 100644 --- a/core/store/src/columns.rs +++ b/core/store/src/columns.rs @@ -293,12 +293,6 @@ pub enum DBCol { /// Witnesses with the lowest index are garbage collected first. /// u64 -> LatestWitnessesKey LatestWitnessesByIndex, - /// Mapping of ShardUId to the underlying ShardUId database key prefix for the State column. - /// It could be parent shard after resharding, ancestor shard in case of many resharding, - /// or just map shard to itself if there was no resharding or we updated the mapping during state sync. - /// - *Rows*: `ShardUId` - /// - *Column type*: `ShardUId` - ShardUIdMapping, } /// Defines different logical parts of a db key. @@ -500,8 +494,7 @@ impl DBCol { | DBCol::FlatState | DBCol::FlatStateChanges | DBCol::FlatStateDeltaMetadata - | DBCol::FlatStorageStatus - | DBCol::ShardUIdMapping => false, + | DBCol::FlatStorageStatus => false, } } @@ -573,7 +566,6 @@ impl DBCol { DBCol::StateTransitionData => &[DBKeyType::BlockHash, DBKeyType::ShardId], DBCol::LatestChunkStateWitnesses => &[DBKeyType::LatestWitnessesKey], DBCol::LatestWitnessesByIndex => &[DBKeyType::LatestWitnessIndex], - DBCol::ShardUIdMapping => &[DBKeyType::ShardUId], } } } diff --git a/core/store/src/flat/inlining_migration.rs b/core/store/src/flat/inlining_migration.rs index 27607e656dc..ecc1f96c557 100644 --- a/core/store/src/flat/inlining_migration.rs +++ b/core/store/src/flat/inlining_migration.rs @@ -17,6 +17,7 @@ use crate::metrics::flat_state_metrics::inlining_migration::{ FLAT_STATE_PAUSED_DURATION, INLINED_COUNT, INLINED_TOTAL_VALUES_SIZE, PROCESSED_COUNT, PROCESSED_TOTAL_VALUES_SIZE, SKIPPED_COUNT, }; +use crate::trie::StateReader; use crate::{DBCol, Store, TrieDBStorage, TrieStorage}; use super::store_helper::{ @@ -96,13 +97,14 @@ struct StateValueReader { } impl StateValueReader { - fn new(store: Store, num_threads: usize) -> Self { + fn new(store: Store, state_reader: StateReader, num_threads: usize) -> Self { let (value_request_send, value_request_recv) = channel::unbounded(); let (value_response_send, value_response_recv) = channel::unbounded(); let mut join_handles = Vec::new(); for _ in 0..num_threads { join_handles.push(Self::spawn_read_value_thread( store.clone(), + state_reader.clone(), value_request_recv.clone(), value_response_send.clone(), )); @@ -130,12 +132,14 @@ impl StateValueReader { fn spawn_read_value_thread( store: Store, + state_reader: StateReader, recv: channel::Receiver, send: channel::Sender, ) -> std::thread::JoinHandle<()> { std::thread::spawn(move || { while let Ok(req) = recv.recv() { - let trie_storage = TrieDBStorage::new(store.clone(), req.shard_uid); + let trie_storage = + TrieDBStorage::new(store.clone(), req.shard_uid, state_reader.clone()); let bytes = match trie_storage.retrieve_raw_bytes(&req.value_hash) { Ok(bytes) => Some(bytes.to_vec()), Err(err) => { @@ -175,7 +179,9 @@ pub fn inline_flat_state_values( ) -> bool { info!(target: "store", %read_state_threads, %batch_size, "Starting FlatState value inlining migration"); let migration_start = std::time::Instant::now(); - let mut value_reader = StateValueReader::new(store.clone(), read_state_threads); + // TODO(reshardingV3) We might want to provide shard_uids together with resharding tree to the `inline_flat_state_values()`. + let state_reader = StateReader::new(&[], [].into()); + let mut value_reader = StateValueReader::new(store.clone(), state_reader, read_state_threads); let mut inlined_total_count = 0; let mut interrupted = false; for (batch_index, batch) in @@ -311,7 +317,7 @@ mod tests { use super::inline_flat_state_values; use crate::flat::store_helper::encode_flat_state_db_key; use crate::flat::{FlatStateValuesInliningMigrationHandle, FlatStorageManager}; - use crate::{DBCol, NodeStorage, Store, TrieCachingStorage}; + use crate::{get_key_from_shard_uid_db_prefix_and_hash, DBCol, NodeStorage, Store}; use borsh::BorshDeserialize; use near_o11y::testonly::init_test_logger; use near_primitives::hash::{hash, CryptoHash}; @@ -439,10 +445,8 @@ mod tests { for (i, value) in values.iter().enumerate() { // TODO(reshardingV3) Test with shard_uid mapping to different shard_uid let shard_uid_db_prefix = shard_uid.into(); - let trie_key = TrieCachingStorage::get_key_from_shard_uid_db_prefix_and_hash( - shard_uid_db_prefix, - &hash(&value), - ); + let trie_key = + get_key_from_shard_uid_db_prefix_and_hash(shard_uid_db_prefix, &hash(&value)); store_update.increment_refcount(DBCol::State, &trie_key, &value); let fs_key = encode_flat_state_db_key(shard_uid, &[i as u8]); let fs_value = borsh::to_vec(&FlatStateValue::value_ref(&value)).unwrap(); diff --git a/core/store/src/flat/store_helper.rs b/core/store/src/flat/store_helper.rs index 3022576e41d..3cc6b6512d4 100644 --- a/core/store/src/flat/store_helper.rs +++ b/core/store/src/flat/store_helper.rs @@ -113,6 +113,10 @@ pub fn remove_delta(store_update: &mut StoreUpdate, shard_uid: ShardUId, block_h /// Must not be called for DBCol::State column, since ReshardingV3 introduces mapping strategy for this column. fn remove_range_by_shard_uid(store_update: &mut StoreUpdate, shard_uid: ShardUId, col: DBCol) { + assert!( + col != DBCol::State, + "Programming error, called remove_range_by_shard_uid() on DBCol::State" + ); let key_from = shard_uid.to_bytes(); let key_to = ShardUId::next_shard_prefix(&key_from); store_update.delete_range(col, &key_from, &key_to); diff --git a/core/store/src/genesis/initialization.rs b/core/store/src/genesis/initialization.rs index ed35943ce95..80439e66cab 100644 --- a/core/store/src/genesis/initialization.rs +++ b/core/store/src/genesis/initialization.rs @@ -17,8 +17,8 @@ use tracing::{error, info, warn}; use crate::{ flat::FlatStorageManager, genesis::GenesisStateApplier, get_genesis_hash, - get_genesis_state_roots, set_genesis_hash, set_genesis_state_roots, ShardTries, - StateSnapshotConfig, Store, TrieConfig, + get_genesis_state_roots, set_genesis_hash, set_genesis_state_roots, trie::StateReader, + ShardTries, StateSnapshotConfig, Store, TrieConfig, }; const STATE_DUMP_FILE: &str = "state_dump"; @@ -111,6 +111,7 @@ fn genesis_state_from_genesis( let runtime_config = runtime_config_store.get_config(genesis.config.protocol_version); let storage_usage_config = &runtime_config.fees.storage_usage_config; let shard_uids: Vec<_> = shard_layout.shard_uids().collect(); + let state_reader = StateReader::new_without_resharding_history(&shard_uids); // note that here we are depending on the behavior that shard_layout.shard_uids() returns an iterator // in order by shard id from 0 to num_shards() let mut shard_account_ids: Vec> = @@ -132,6 +133,7 @@ fn genesis_state_from_genesis( store.clone(), TrieConfig::default(), &shard_uids, + state_reader, FlatStorageManager::new(store), StateSnapshotConfig::default(), ); diff --git a/core/store/src/lib.rs b/core/store/src/lib.rs index 920d6285de5..9416716ed10 100644 --- a/core/store/src/lib.rs +++ b/core/store/src/lib.rs @@ -6,10 +6,10 @@ use crate::db::{refcount, DBIterator, DBOp, DBSlice, DBTransaction, Database, St pub use crate::trie::iterator::{TrieIterator, TrieTraversalItem}; pub use crate::trie::update::{TrieUpdate, TrieUpdateIterator, TrieUpdateValuePtr}; pub use crate::trie::{ - estimator, map_shard_uid, resharding_v2, ApplyStatePartResult, KeyForStateChanges, - KeyLookupMode, NibbleSlice, PartialStorage, PrefetchApi, PrefetchError, RawTrieNode, - RawTrieNodeWithSize, ShardTries, StateSnapshot, StateSnapshotConfig, Trie, TrieAccess, - TrieCache, TrieCachingStorage, TrieChanges, TrieConfig, TrieDBStorage, TrieStorage, + estimator, get_key_from_shard_uid_db_prefix_and_hash, resharding_v2, ApplyStatePartResult, + KeyForStateChanges, KeyLookupMode, NibbleSlice, PartialStorage, PrefetchApi, PrefetchError, + RawTrieNode, RawTrieNodeWithSize, ShardTries, StateSnapshot, StateSnapshotConfig, Trie, + TrieAccess, TrieCache, TrieCachingStorage, TrieChanges, TrieConfig, TrieDBStorage, TrieStorage, WrappedTrieChanges, STATE_SNAPSHOT_COLUMNS, }; use borsh::{BorshDeserialize, BorshSerialize}; diff --git a/core/store/src/test_utils.rs b/core/store/src/test_utils.rs index 1c66a7b29a2..0e758d4220e 100644 --- a/core/store/src/test_utils.rs +++ b/core/store/src/test_utils.rs @@ -3,6 +3,7 @@ use crate::flat::{ store_helper, BlockInfo, FlatStorageManager, FlatStorageReadyStatus, FlatStorageStatus, }; use crate::metadata::{DbKind, DbVersion, DB_VERSION}; +use crate::trie::StateReader; use crate::{ get, get_delayed_receipt_indices, get_promise_yield_indices, DBCol, NodeStorage, ShardTries, StateSnapshotConfig, Store, Trie, TrieConfig, @@ -63,6 +64,10 @@ pub fn create_test_node_storage_with_cold( (storage, hot, cold) } +pub fn create_test_state_reader(shard_uid: ShardUId) -> StateReader { + StateReader::new_without_resharding_history(&[shard_uid]) +} + /// Creates an in-memory database. pub fn create_test_store() -> Store { create_test_node_storage(DB_VERSION, DbKind::RPC).get_hot_store() @@ -123,6 +128,7 @@ impl TestTriesBuilder { let shard_uids = (0..self.num_shards) .map(|shard_id| ShardUId { shard_id: shard_id as u32, version: self.shard_version }) .collect::>(); + let state_reader = StateReader::new_without_resharding_history(&shard_uids); let flat_storage_manager = FlatStorageManager::new(store.clone()); let tries = ShardTries::new( store.clone(), @@ -131,6 +137,7 @@ impl TestTriesBuilder { ..Default::default() }, &shard_uids, + state_reader, flat_storage_manager, StateSnapshotConfig::default(), ); diff --git a/core/store/src/trie/from_flat.rs b/core/store/src/trie/from_flat.rs index 71cf52684ea..6c1d7a387d1 100644 --- a/core/store/src/trie/from_flat.rs +++ b/core/store/src/trie/from_flat.rs @@ -1,4 +1,5 @@ use crate::flat::{store_helper, FlatStorageError, FlatStorageManager}; +use crate::trie::StateReader; use crate::{ShardTries, StateSnapshotConfig, Store, Trie, TrieConfig, TrieDBStorage, TrieStorage}; use near_primitives::{shard_layout::ShardUId, state::FlatStateValue}; use std::time::Instant; @@ -11,7 +12,9 @@ use std::time::Instant; // Please note that the trie is created for the block state with height equal to flat_head // flat state can comtain deltas after flat_head and can be different from tip of the blockchain. pub fn construct_trie_from_flat(store: Store, write_store: Store, shard_uid: ShardUId) { - let trie_storage = TrieDBStorage::new(store.clone(), shard_uid); + let shard_uids = &[shard_uid]; + let state_reader = StateReader::new_without_resharding_history(shard_uids); + let trie_storage = TrieDBStorage::new(store.clone(), shard_uid, state_reader.clone()); let flat_state_to_trie_kv = |entry: Result<(Vec, FlatStateValue), FlatStorageError>| -> (Vec, Vec) { let (key, value) = entry.unwrap(); @@ -31,7 +34,8 @@ pub fn construct_trie_from_flat(store: Store, write_store: Store, shard_uid: Sha let tries = ShardTries::new( write_store.clone(), TrieConfig::default(), - &[shard_uid], + shard_uids, + state_reader, FlatStorageManager::new(write_store), StateSnapshotConfig::default(), ); diff --git a/core/store/src/trie/mem/parallel_loader.rs b/core/store/src/trie/mem/parallel_loader.rs index 7fdcfdb6534..11d4f21d542 100644 --- a/core/store/src/trie/mem/parallel_loader.rs +++ b/core/store/src/trie/mem/parallel_loader.rs @@ -5,11 +5,11 @@ use super::construction::TrieConstructor; use super::node::{InputMemTrieNode, MemTrieNodeId}; use crate::flat::FlatStorageError; use crate::trie::Children; -use crate::{map_shard_uid, DBCol, NibbleSlice, RawTrieNode, RawTrieNodeWithSize, Store}; +use crate::{DBCol, NibbleSlice, RawTrieNode, RawTrieNodeWithSize, Store}; use borsh::BorshDeserialize; use near_primitives::errors::{MissingTrieValueContext, StorageError}; use near_primitives::hash::CryptoHash; -use near_primitives::shard_layout::{MappedShardUId, ShardUId}; +use near_primitives::shard_layout::ShardUId; use near_primitives::state::FlatStateValue; use near_primitives::types::StateRoot; use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator}; @@ -46,7 +46,7 @@ pub fn load_memtrie_in_parallel( /// or multiple shards. pub struct ParallelMemTrieLoader { store: Store, - shard_uid_db_prefix: MappedShardUId, + shard_uid: ShardUId, root: StateRoot, num_subtrees_desired: usize, } @@ -58,8 +58,7 @@ impl ParallelMemTrieLoader { root: StateRoot, num_subtrees_desired: usize, ) -> Self { - let shard_uid_db_prefix = map_shard_uid(&store, shard_uid); - Self { store, shard_uid_db_prefix, root, num_subtrees_desired } + Self { store, shard_uid, root, num_subtrees_desired } } /// Implements stage 1; recursively expanding the trie until all subtrees are small enough. @@ -89,7 +88,8 @@ impl ParallelMemTrieLoader { ) -> Result { // Read the node from the State column. let mut key = [0u8; 40]; - key[0..8].copy_from_slice(&self.shard_uid_db_prefix.0.to_bytes()); + // TODO(reshardingV3) Likely should use mapped shard UId here + key[0..8].copy_from_slice(&self.shard_uid.to_bytes()); key[8..40].copy_from_slice(&hash.0); let node = RawTrieNodeWithSize::try_from_slice( &self @@ -215,7 +215,8 @@ impl ParallelMemTrieLoader { arena: &mut impl ArenaMut, ) -> Result { // Figure out which range corresponds to the prefix of this subtree. - let (start, end) = subtree_to_load.to_iter_range(self.shard_uid_db_prefix.0); + // TODO(reshardingV3) likely should use mapped shard UId here + let (start, end) = subtree_to_load.to_iter_range(self.shard_uid); // Load all the keys in this range from the FlatState column. let mut recon = TrieConstructor::new(arena); diff --git a/core/store/src/trie/mod.rs b/core/store/src/trie/mod.rs index e790121eeb7..9affa8c349b 100644 --- a/core/store/src/trie/mod.rs +++ b/core/store/src/trie/mod.rs @@ -14,12 +14,11 @@ use crate::trie::iterator::TrieIterator; pub use crate::trie::nibble_slice::NibbleSlice; pub use crate::trie::prefetching_trie_storage::{PrefetchApi, PrefetchError}; pub use crate::trie::shard_tries::{KeyForStateChanges, ShardTries, WrappedTrieChanges}; +pub use crate::trie::state_reader::{get_key_from_shard_uid_db_prefix_and_hash, StateReader}; pub use crate::trie::state_snapshot::{ SnapshotError, StateSnapshot, StateSnapshotConfig, STATE_SNAPSHOT_COLUMNS, }; -pub use crate::trie::trie_storage::{ - map_shard_uid, TrieCache, TrieCachingStorage, TrieDBStorage, TrieStorage, -}; +pub use crate::trie::trie_storage::{TrieCache, TrieCachingStorage, TrieDBStorage, TrieStorage}; use crate::StorageError; use borsh::{BorshDeserialize, BorshSerialize}; pub use from_flat::construct_trie_from_flat; @@ -56,6 +55,7 @@ pub mod receipts_column_helper; pub mod resharding_v2; mod shard_tries; mod state_parts; +mod state_reader; mod state_snapshot; mod trie_recording; mod trie_storage; diff --git a/core/store/src/trie/prefetching_trie_storage.rs b/core/store/src/trie/prefetching_trie_storage.rs index 0fc112f682f..1a9f468e3d1 100644 --- a/core/store/src/trie/prefetching_trie_storage.rs +++ b/core/store/src/trie/prefetching_trie_storage.rs @@ -1,22 +1,23 @@ use crate::config::PrefetchConfig; -use crate::map_shard_uid; use crate::sync_utils::Monitor; use crate::{ - metrics, DBCol, MissingTrieValueContext, StorageError, Store, Trie, TrieCache, - TrieCachingStorage, TrieConfig, TrieStorage, + metrics, DBCol, MissingTrieValueContext, StorageError, Store, Trie, TrieCache, TrieConfig, + TrieStorage, }; use crossbeam::select; use near_o11y::metrics::prometheus; use near_o11y::metrics::prometheus::core::GenericGauge; use near_o11y::tracing::error; use near_primitives::hash::CryptoHash; -use near_primitives::shard_layout::{MappedShardUId, ShardUId}; +use near_primitives::shard_layout::ShardUId; use near_primitives::trie_key::TrieKey; use near_primitives::types::{AccountId, ShardId, StateRoot}; use std::collections::HashMap; use std::sync::Arc; use std::thread; +use super::StateReader; + const MAX_QUEUED_WORK_ITEMS: usize = 16 * 1024; const MAX_PREFETCH_STAGING_MEMORY: usize = 200 * 1024 * 1024; /// How much memory capacity is reserved for each prefetch request before @@ -49,12 +50,13 @@ const NUM_IO_THREADS: usize = 8; struct TriePrefetchingStorage { /// Store is shared with parent `TrieCachingStorage`. store: Store, - shard_uid_db_prefix: MappedShardUId, + shard_uid: ShardUId, /// Shard cache is shared with parent `TrieCachingStorage`. But the /// pre-fetcher uses this in read-only mode to avoid premature evictions. shard_cache: TrieCache, /// Shared with parent `TrieCachingStorage`. prefetching: PrefetchStagingArea, + state_reader: StateReader, } /// This type is shared between runtime crate and store crate. @@ -80,6 +82,7 @@ pub struct PrefetchApi { store: Store, shard_cache: TrieCache, + state_reader: StateReader, pub enable_receipt_prefetching: bool, /// Configured accounts will be prefetched as SWEAT token account, if predecessor is listed as receiver. @@ -244,10 +247,7 @@ impl TrieStorage for TriePrefetchingStorage { match prefetch_state { // Slot reserved for us, this thread should fetch it from DB. PrefetcherResult::SlotReserved => { - let key = TrieCachingStorage::get_key_from_shard_uid_db_prefix_and_hash( - self.shard_uid_db_prefix, - hash, - ); + let key = self.state_reader.get_key_from_shard_uid_and_hash(self.shard_uid, hash); match self.store.get(DBCol::State, key.as_ref()) { Ok(Some(value)) => { let value: Arc<[u8]> = value.into(); @@ -311,9 +311,9 @@ impl TriePrefetchingStorage { shard_uid: ShardUId, shard_cache: TrieCache, prefetching: PrefetchStagingArea, + state_reader: StateReader, ) -> Self { - let shard_uid_db_prefix = map_shard_uid(&store, shard_uid); - Self { store, shard_uid_db_prefix, shard_cache, prefetching } + Self { store, shard_uid, shard_cache, prefetching, state_reader } } } @@ -409,6 +409,7 @@ impl PrefetchApi { pub(crate) fn new( store: Store, shard_cache: TrieCache, + state_reader: StateReader, shard_uid: ShardUId, trie_config: &TrieConfig, ) -> (Self, PrefetchingThreadsHandle) { @@ -430,6 +431,7 @@ impl PrefetchApi { shard_uid, store, shard_cache, + state_reader, }; let (shutdown_tx, shutdown_rx) = crossbeam::channel::bounded(1); let handles = (0..NUM_IO_THREADS) @@ -456,6 +458,7 @@ impl PrefetchApi { self.shard_uid, self.shard_cache.clone(), self.prefetching.clone(), + self.state_reader.clone(), )) } @@ -469,6 +472,7 @@ impl PrefetchApi { self.shard_uid, self.shard_cache.clone(), self.prefetching.clone(), + self.state_reader.clone(), ); let work_queue = self.work_queue_rx.clone(); let metric_prefetch_sent = diff --git a/core/store/src/trie/shard_tries.rs b/core/store/src/trie/shard_tries.rs index b7ee48d4b67..47b19f22906 100644 --- a/core/store/src/trie/shard_tries.rs +++ b/core/store/src/trie/shard_tries.rs @@ -1,8 +1,7 @@ use super::mem::mem_tries::MemTries; use super::state_snapshot::{StateSnapshot, StateSnapshotConfig}; -use super::TrieRefcountSubtraction; +use super::{StateReader, TrieRefcountSubtraction}; use crate::flat::{FlatStorageManager, FlatStorageStatus}; -use crate::map_shard_uid; use crate::trie::config::TrieConfig; use crate::trie::mem::loading::load_trie_from_flat_state_and_delta; use crate::trie::prefetching_trie_storage::PrefetchingThreadsHandle; @@ -12,7 +11,7 @@ use crate::{metrics, DBCol, PrefetchApi, TrieDBStorage, TrieStorage}; use crate::{Store, StoreUpdate, Trie, TrieChanges, TrieUpdate}; use near_primitives::errors::StorageError; use near_primitives::hash::CryptoHash; -use near_primitives::shard_layout::{self, MappedShardUId, ShardUId}; +use near_primitives::shard_layout::{self, ShardUId}; use near_primitives::trie_key::TrieKey; use near_primitives::types::{ BlockHeight, RawStateChange, RawStateChangesWithTrieKey, StateChangeCause, StateRoot, @@ -40,8 +39,8 @@ struct ShardTriesInner { state_snapshot: Arc>>, /// Configures how to make state snapshots. state_snapshot_config: StateSnapshotConfig, - /// Tells which ShardUId prefix should be used instead when accessing db State column. - shard_uid_to_db_prefix: HashMap, + /// Wrapper to access the db State column. + state_reader: StateReader, } #[derive(Clone)] @@ -52,15 +51,12 @@ impl ShardTries { store: Store, trie_config: TrieConfig, shard_uids: &[ShardUId], + state_reader: StateReader, flat_storage_manager: FlatStorageManager, state_snapshot_config: StateSnapshotConfig, ) -> Self { let caches = Self::create_initial_caches(&trie_config, &shard_uids, false); let view_caches = Self::create_initial_caches(&trie_config, &shard_uids, true); - let shard_uid_to_db_prefix: HashMap = shard_uids - .iter() - .map(|&shard_uid| (shard_uid, map_shard_uid(&store, shard_uid))) - .collect(); metrics::HAS_STATE_SNAPSHOT.set(0); ShardTries(Arc::new(ShardTriesInner { store, @@ -72,7 +68,7 @@ impl ShardTries { prefetchers: Default::default(), state_snapshot: Arc::new(RwLock::new(None)), state_snapshot_config, - shard_uid_to_db_prefix, + state_reader, })) } @@ -129,7 +125,11 @@ impl ShardTries { if let Some(cache) = self.get_trie_cache_for(shard_uid, is_view) { Arc::new(self.create_caching_storage(cache, shard_uid, is_view)) } else { - Arc::new(TrieDBStorage::new(self.0.store.clone(), shard_uid)) + Arc::new(TrieDBStorage::new( + self.0.store.clone(), + shard_uid, + self.0.state_reader.clone(), + )) }; let flat_storage_chunk_view = block_hash .and_then(|block_hash| self.0.flat_storage_manager.chunk_view(shard_uid, block_hash)); @@ -153,7 +153,14 @@ impl ShardTries { let cache = self .get_trie_cache_for(shard_uid, true) .expect("trie cache should be enabled for view calls"); - let storage = Arc::new(TrieCachingStorage::new(store, cache, shard_uid, true, None)); + let storage = Arc::new(TrieCachingStorage::new( + store, + cache, + shard_uid, + true, + None, + self.0.state_reader.clone(), + )); let flat_storage_chunk_view = flat_storage_manager.chunk_view(shard_uid, *block_hash); Ok(Trie::new(storage, state_root, flat_storage_chunk_view)) @@ -181,6 +188,10 @@ impl ShardTries { self.0.store.clone() } + pub fn get_state_reader(&self) -> StateReader { + self.0.state_reader.clone() + } + pub(crate) fn get_db(&self) -> &Arc { &self.0.store.storage } @@ -232,6 +243,7 @@ impl ShardTries { PrefetchApi::new( self.0.store.clone(), cache.clone(), + self.0.state_reader.clone(), shard_uid, &self.0.trie_config, ) @@ -239,7 +251,14 @@ impl ShardTries { .0 .clone() }); - TrieCachingStorage::new(self.0.store.clone(), cache, shard_uid, is_view, prefetch_api) + TrieCachingStorage::new( + self.0.store.clone(), + cache, + shard_uid, + is_view, + prefetch_api, + self.0.state_reader.clone(), + ) } fn apply_deletions_inner( @@ -250,15 +269,10 @@ impl ShardTries { ) { let mut ops = Vec::with_capacity(deletions.len()); for TrieRefcountSubtraction { trie_node_or_value_hash, rc, .. } in deletions.iter() { - let shard_uid_db_prefix = match self.0.shard_uid_to_db_prefix.get(&shard_uid) { - Some(mapped_shard_uid) => *mapped_shard_uid, - // TODO(reshardingV3) Think about how None should be handled here. - None => shard_uid.into(), - }; - let key = TrieCachingStorage::get_key_from_shard_uid_db_prefix_and_hash( - shard_uid_db_prefix, - trie_node_or_value_hash, - ); + let key = self + .0 + .state_reader + .get_key_from_shard_uid_and_hash(shard_uid, trie_node_or_value_hash); store_update.decrement_refcount_by(DBCol::State, key.as_ref(), *rc); ops.push((trie_node_or_value_hash, None)); } @@ -276,15 +290,10 @@ impl ShardTries { for TrieRefcountAddition { trie_node_or_value_hash, trie_node_or_value, rc } in insertions.iter() { - let shard_uid_db_prefix = match self.0.shard_uid_to_db_prefix.get(&shard_uid) { - Some(mapped_shard_uid) => *mapped_shard_uid, - // TODO(reshardingV3) Think about how None should be handled here. - None => shard_uid.into(), - }; - let key = TrieCachingStorage::get_key_from_shard_uid_db_prefix_and_hash( - shard_uid_db_prefix, - trie_node_or_value_hash, - ); + let key = self + .0 + .state_reader + .get_key_from_shard_uid_and_hash(shard_uid, trie_node_or_value_hash); store_update.increment_refcount_by(DBCol::State, key.as_ref(), trie_node_or_value, *rc); ops.push((trie_node_or_value_hash, Some(trie_node_or_value.as_slice()))); } @@ -771,10 +780,12 @@ mod test { ..TrieConfig::default() }; let shard_uids = Vec::from([ShardUId::single_shard()]); + let state_reader = StateReader::new_without_resharding_history(&shard_uids); ShardTries::new( store.clone(), trie_config, &shard_uids, + state_reader, FlatStorageManager::new(store), StateSnapshotConfig::default(), ) @@ -888,11 +899,13 @@ mod test { }; let shard_uids = Vec::from([ShardUId { shard_id: 0, version: 0 }]); let shard_uid = *shard_uids.first().unwrap(); + let state_reader = StateReader::new_without_resharding_history(&shard_uids); let trie = ShardTries::new( store.clone(), trie_config, &shard_uids, + state_reader, FlatStorageManager::new(store), StateSnapshotConfig::default(), ); diff --git a/core/store/src/trie/state_parts.rs b/core/store/src/trie/state_parts.rs index c5d502ec160..f2e0f05bde3 100644 --- a/core/store/src/trie/state_parts.rs +++ b/core/store/src/trie/state_parts.rs @@ -525,7 +525,7 @@ mod tests { }; use super::*; - use crate::{DBCol, MissingTrieValueContext, TrieCachingStorage}; + use crate::{get_key_from_shard_uid_db_prefix_and_hash, DBCol, MissingTrieValueContext}; use near_primitives::shard_layout::ShardUId; /// Checks that sampling state boundaries always gives valid state keys @@ -1230,10 +1230,7 @@ mod tests { let value_hash = hash(&store_value); // TODO(reshardingV3) Test shard UId mapping to different shard UId. let shard_uid_db_prefix = shard_uid.into(); - let store_key = TrieCachingStorage::get_key_from_shard_uid_db_prefix_and_hash( - shard_uid_db_prefix, - &value_hash, - ); + let store_key = get_key_from_shard_uid_db_prefix_and_hash(shard_uid_db_prefix, &value_hash); store_update.decrement_refcount(DBCol::State, &store_key); store_update.commit().unwrap(); diff --git a/core/store/src/trie/state_reader.rs b/core/store/src/trie/state_reader.rs new file mode 100644 index 00000000000..77dfec60930 --- /dev/null +++ b/core/store/src/trie/state_reader.rs @@ -0,0 +1,153 @@ +use crate::trie::POISONED_LOCK_ERR; +use crate::{DBCol, MissingTrieValueContext, StorageError, Store}; +use near_primitives::hash::CryptoHash; +use near_primitives::shard_layout::ShardUId; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +/// This file implements structs and helper functions needed for +/// the mapping strategy for State column in ReshardingV3. + +/// A wrapper for ShardUId used to denote when we expect the mapped version of the value. +/// +/// ReshardingV3 uses mapping strategy for the State column. +/// It means we access trie storage using a mapped shard UID, that can be the same shard UID, +/// resharding parent shard UID, or any ancestor shard UID in the tree of resharding history. +#[derive(Clone, Copy, PartialEq)] +pub struct MappedShardUId(pub ShardUId); + +impl From for MappedShardUId { + fn from(shard_uid: ShardUId) -> Self { + MappedShardUId(shard_uid) + } +} + +// TODO(reshardingV3) Go through all places where this method is used. +// and see whether the provided `shard_uid_db_prefix` is correct. +// Perhaps we would want to make this function private and only use +// `StateReader::get_key_from_shard_uid_and_hash()`. +pub fn get_key_from_shard_uid_db_prefix_and_hash( + shard_uid_db_prefix: MappedShardUId, + hash: &CryptoHash, +) -> [u8; 40] { + let mut key = [0; 40]; + key[0..8].copy_from_slice(&shard_uid_db_prefix.0.to_bytes()); + key[8..].copy_from_slice(hash.as_ref()); + key +} + +/// A wrapper to read from State column (Trie storage). +/// It maintains in-memory mapping from ShardUId to a db key shard_uid prefix +/// that should be used when reading trie node from storage. +/// +/// We protect internal state with Mutex for concurrent use. +#[derive(Clone)] +pub struct StateReader(Arc>); + +pub struct StateReaderInner { + /// Tells which ShardUId prefix should be used instead when accessing db State column. + shard_uid_to_db_prefix: HashMap, + /// Stores resharding tree (mapping from shard_uid to parent shard_uid). + shard_uid_to_parent: HashMap>, +} + +impl StateReaderInner { + pub(crate) fn get_mapped_shard_uid(&mut self, shard_uid: ShardUId) -> MappedShardUId { + match self.shard_uid_to_db_prefix.get(&shard_uid) { + Some(ancestor) => *ancestor, + None => { + // TODO(reshardingV3) That means shard UId appeared for which + // we did not explicitly set the mapping. Think how to correctly handle this. + self.shard_uid_to_db_prefix.insert(shard_uid, shard_uid.into()); + shard_uid.into() + } + } + } + + pub(crate) fn get_parent_shard_uid(&self, shard_uid: ShardUId) -> Option { + match self.shard_uid_to_parent.get(&shard_uid) { + Some(maybe_parent_shard_uid) => *maybe_parent_shard_uid, + None => { + // TODO(reshardingV3) That means inconsistent state. Each shard should have + // its parent set (None for root shard). Return an error here. + None + } + } + } + + pub(crate) fn update_shard_uid_mapping( + &mut self, + shard_uid: ShardUId, + mapped_shard_uid: MappedShardUId, + ) { + self.shard_uid_to_db_prefix.insert(shard_uid, mapped_shard_uid); + } +} + +impl StateReader { + pub fn new( + shard_uids: &[ShardUId], + shard_uid_to_parent: HashMap>, + ) -> Self { + let shard_uid_to_db_prefix: HashMap = + shard_uids.iter().map(|&shard_uid| (shard_uid, shard_uid.into())).collect(); + Self(Arc::new(Mutex::new(StateReaderInner { shard_uid_to_db_prefix, shard_uid_to_parent }))) + } + + // TODO(reshardingV3) Go through all places where this method is used + // and see whether it should be replaced with `new()` with real resharding tree as an argument. + pub fn new_without_resharding_history(shard_uids: &[ShardUId]) -> Self { + let shard_uid_to_parent = shard_uids.iter().map(|id| (*id, None)).collect(); + Self::new(shard_uids, shard_uid_to_parent) + } + + // TODO(reshardingV3) This does not update the mapping if the key does not exist. + // Look at places where this method is used and think about what to do about this. + pub fn get_key_from_shard_uid_and_hash( + &self, + shard_uid: ShardUId, + hash: &CryptoHash, + ) -> [u8; 40] { + let mapped_shard_uid = self.lock().get_mapped_shard_uid(shard_uid); + get_key_from_shard_uid_db_prefix_and_hash(mapped_shard_uid, hash) + } + + /// Main function of `StateReader`. + /// Attempts to read Trie node/value using the in-memory shard_uid mapping for db key prefix. + /// It recursively uses parent shard_uid if the key was not found in the database. + /// Eventually, if the read succeded with ancestor shard_uid, the in-memory mapping is updated + /// so that the next time we attempt to read a node for this shard we start with the correct mapping. + pub fn read( + &self, + store: &Store, + shard_uid: ShardUId, + hash: &CryptoHash, + ) -> Result, StorageError> { + let initial_mapped_shard_uid = self.lock().get_mapped_shard_uid(shard_uid); + let mut mapped_shard_uid = initial_mapped_shard_uid; + let result = loop { + let key = get_key_from_shard_uid_db_prefix_and_hash(mapped_shard_uid, hash); + let val = store + .get(DBCol::State, key.as_ref()) + .map_err(|_| StorageError::StorageInternalError)?; + if let Some(val) = val { + break Some(val.into()); + } + mapped_shard_uid = match self.lock().get_parent_shard_uid(mapped_shard_uid.0) { + Some(parent_shard_uid) => parent_shard_uid.into(), + None => break None, + }; + }; + + // Only update the mapping if the trie node/value was finally found in the storage. + if result.is_some() && mapped_shard_uid != initial_mapped_shard_uid { + self.lock().update_shard_uid_mapping(shard_uid, mapped_shard_uid); + } + + result.ok_or(StorageError::MissingTrieValue(MissingTrieValueContext::TrieStorage, *hash)) + } + + fn lock(&self) -> std::sync::MutexGuard { + self.0.lock().expect(POISONED_LOCK_ERR) + } +} diff --git a/core/store/src/trie/trie_storage.rs b/core/store/src/trie/trie_storage.rs index 0c448180c16..73aed271a9d 100644 --- a/core/store/src/trie/trie_storage.rs +++ b/core/store/src/trie/trie_storage.rs @@ -1,19 +1,21 @@ use crate::trie::config::TrieConfig; use crate::trie::prefetching_trie_storage::PrefetcherResult; use crate::trie::POISONED_LOCK_ERR; -use crate::{metrics, DBCol, MissingTrieValueContext, PrefetchApi, StorageError, Store}; +use crate::{metrics, MissingTrieValueContext, PrefetchApi, StorageError, Store}; use lru::LruCache; use near_o11y::log_assert; use near_o11y::metrics::prometheus; use near_o11y::metrics::prometheus::core::{GenericCounter, GenericGauge}; use near_primitives::challenge::PartialState; use near_primitives::hash::CryptoHash; -use near_primitives::shard_layout::{MappedShardUId, ShardUId}; +use near_primitives::shard_layout::ShardUId; use near_primitives::types::ShardId; use std::collections::{HashMap, HashSet, VecDeque}; use std::num::NonZeroUsize; use std::sync::{Arc, Mutex}; +use super::StateReader; + pub(crate) struct BoundedQueue { queue: VecDeque, /// If queue size exceeds capacity, item from the tail is removed. @@ -362,8 +364,6 @@ pub struct TrieCachingStorage { pub(crate) shard_uid: ShardUId, pub(crate) is_view: bool, - pub(crate) shard_uid_db_prefix: MappedShardUId, - /// Caches ever requested items for the shard `shard_uid`. Used to speed up DB operations, presence of any item is /// not guaranteed. pub(crate) shard_cache: TrieCache, @@ -371,6 +371,8 @@ pub struct TrieCachingStorage { /// The entry point for the runtime to submit prefetch requests. pub(crate) prefetch_api: Option, + state_reader: StateReader, + // Counters tracking operations happening inside the shard cache. // Stored here to avoid overhead of looking them up on hot paths. metrics: TrieCacheInnerMetrics, @@ -397,6 +399,7 @@ impl TrieCachingStorage { shard_uid: ShardUId, is_view: bool, prefetch_api: Option, + state_reader: StateReader, ) -> TrieCachingStorage { // `itoa` is much faster for printing shard_id to a string than trivial alternatives. let mut buffer = itoa::Buffer::new(); @@ -420,28 +423,17 @@ impl TrieCachingStorage { prefetch_retry: metrics::PREFETCH_RETRY.with_label_values(&metrics_labels[..1]), prefetch_conflict: metrics::PREFETCH_CONFLICT.with_label_values(&metrics_labels[..1]), }; - let shard_uid_db_prefix = map_shard_uid(&store, shard_uid); TrieCachingStorage { store, shard_uid, is_view, - shard_uid_db_prefix, shard_cache, prefetch_api, + state_reader, metrics, } } - pub fn get_key_from_shard_uid_db_prefix_and_hash( - shard_uid_db_prefix: MappedShardUId, - hash: &CryptoHash, - ) -> [u8; 40] { - let mut key = [0; 40]; - key[0..8].copy_from_slice(&shard_uid_db_prefix.0.to_bytes()); - key[8..].copy_from_slice(hash.as_ref()); - key - } - /// Reads value if it is not in shard cache. Handles dropping the cache /// lock. Either waits for prefetcher to fetch it or reads it from DB. /// It is responsibility of caller to release the prefetch slot later. @@ -562,33 +554,9 @@ impl TrieStorage for TrieCachingStorage { } } -/// This method reads the shard UID mapping from db. -/// See `MappedShardUID` documentation for explanation. -pub fn map_shard_uid(store: &Store, shard_uid: ShardUId) -> MappedShardUId { - match store.get_ser::(DBCol::ShardUIdMapping, &shard_uid.to_bytes()) { - Ok(Some(mapped_shard_uid)) => mapped_shard_uid, - // TODO(reshardingV3) Think about how to handle it the best way - _ => shard_uid, - } - .into() -} - -fn read_node_from_db( - store: &Store, - shard_uid: MappedShardUId, - hash: &CryptoHash, -) -> Result, StorageError> { - let key = TrieCachingStorage::get_key_from_shard_uid_db_prefix_and_hash(shard_uid, hash); - let val = store - .get(DBCol::State, key.as_ref()) - .map_err(|_| StorageError::StorageInternalError)? - .ok_or(StorageError::MissingTrieValue(MissingTrieValueContext::TrieStorage, *hash))?; - Ok(val.into()) -} - impl TrieCachingStorage { fn read_from_db(&self, hash: &CryptoHash) -> Result, StorageError> { - read_node_from_db(&self.store, self.shard_uid_db_prefix, hash) + self.state_reader.read(&self.store, self.shard_uid, hash) } pub fn prefetch_api(&self) -> &Option { @@ -602,19 +570,19 @@ impl TrieCachingStorage { /// It is useful for background tasks that should not affect chunk processing and block each other. pub struct TrieDBStorage { pub(crate) store: Store, - pub(crate) shard_uid_db_prefix: MappedShardUId, + pub(crate) shard_uid: ShardUId, + state_reader: StateReader, } impl TrieDBStorage { - pub fn new(store: Store, shard_uid: ShardUId) -> Self { - let shard_uid_db_prefix = map_shard_uid(&store, shard_uid); - Self { store, shard_uid_db_prefix } + pub fn new(store: Store, shard_uid: ShardUId, state_reader: StateReader) -> Self { + Self { store, shard_uid, state_reader } } } impl TrieStorage for TrieDBStorage { fn retrieve_raw_bytes(&self, hash: &CryptoHash) -> Result, StorageError> { - read_node_from_db(&self.store, self.shard_uid_db_prefix, hash) + self.state_reader.read(&self.store, self.shard_uid, hash) } } diff --git a/core/store/src/trie/trie_tests.rs b/core/store/src/trie/trie_tests.rs index 68b23311189..c9f05c4aba1 100644 --- a/core/store/src/trie/trie_tests.rs +++ b/core/store/src/trie/trie_tests.rs @@ -203,7 +203,7 @@ mod nodes_counter_tests { #[cfg(test)] mod trie_storage_tests { use super::*; - use crate::test_utils::create_test_store; + use crate::test_utils::{create_test_state_reader, create_test_store}; use crate::trie::accounting_cache::TrieAccountingCache; use crate::trie::trie_storage::{TrieCache, TrieCachingStorage, TrieDBStorage}; use crate::trie::TrieRefcountAddition; @@ -236,7 +236,8 @@ mod trie_storage_tests { let values = vec![value.clone()]; let shard_uid = ShardUId::single_shard(); let store = create_store_with_values(&values, shard_uid); - let trie_db_storage = TrieDBStorage::new(store, shard_uid); + let state_reader = create_test_state_reader(shard_uid); + let trie_db_storage = TrieDBStorage::new(store, shard_uid, state_reader); let key = hash(&value); assert_eq!(trie_db_storage.retrieve_raw_bytes(&key).unwrap().as_ref(), value); let wrong_key = hash(&[2]); @@ -250,9 +251,16 @@ mod trie_storage_tests { let values = vec![value.clone()]; let shard_uid = ShardUId::single_shard(); let store = create_store_with_values(&values, shard_uid); + let state_reader = create_test_state_reader(shard_uid); let trie_cache = TrieCache::new(&TrieConfig::default(), shard_uid, false); - let trie_caching_storage = - TrieCachingStorage::new(store, trie_cache.clone(), shard_uid, false, None); + let trie_caching_storage = TrieCachingStorage::new( + store, + trie_cache.clone(), + shard_uid, + false, + None, + state_reader, + ); let mut accounting_cache = TrieAccountingCache::new(None); let key = hash(&value); assert_eq!(trie_cache.get(&key), None); @@ -275,12 +283,15 @@ mod trie_storage_tests { fn test_retrieve_error() { let shard_uid = ShardUId::single_shard(); let store = create_test_store(); + let state_reader = create_test_state_reader(shard_uid); + let trie_caching_storage = TrieCachingStorage::new( store, TrieCache::new(&TrieConfig::default(), shard_uid, false), shard_uid, false, None, + state_reader, ); let value = vec![1u8]; let key = hash(&value); @@ -296,9 +307,16 @@ mod trie_storage_tests { let values = vec![value.clone()]; let shard_uid = ShardUId::single_shard(); let store = create_store_with_values(&values, shard_uid); + let state_reader = create_test_state_reader(shard_uid); let trie_cache = TrieCache::new(&TrieConfig::default(), shard_uid, false); - let trie_caching_storage = - TrieCachingStorage::new(store, trie_cache.clone(), shard_uid, false, None); + let trie_caching_storage = TrieCachingStorage::new( + store, + trie_cache.clone(), + shard_uid, + false, + None, + state_reader, + ); let mut accounting_cache = TrieAccountingCache::new(None); let key = hash(&value); @@ -322,9 +340,16 @@ mod trie_storage_tests { let values = vec![vec![1u8]]; let shard_uid = ShardUId::single_shard(); let store = create_store_with_values(&values, shard_uid); + let state_reader = create_test_state_reader(shard_uid); let trie_cache = TrieCache::new(&TrieConfig::default(), shard_uid, false); - let trie_caching_storage = - TrieCachingStorage::new(store, trie_cache.clone(), shard_uid, false, None); + let trie_caching_storage = TrieCachingStorage::new( + store, + trie_cache.clone(), + shard_uid, + false, + None, + state_reader, + ); let mut accounting_cache = TrieAccountingCache::new(None); let value = &values[0]; let key = hash(&value); @@ -379,14 +404,21 @@ mod trie_storage_tests { let values: Vec> = (0..shard_cache_size as u8 + 1).map(|i| vec![i]).collect(); let shard_uid = ShardUId::single_shard(); let store = create_store_with_values(&values, shard_uid); + let state_reader = create_test_state_reader(shard_uid); let mut trie_config = TrieConfig::default(); trie_config .shard_cache_config .per_shard_max_bytes .insert(shard_uid, bytesize::ByteSize(shard_cache_size)); let trie_cache = TrieCache::new(&trie_config, shard_uid, false); - let trie_caching_storage = - TrieCachingStorage::new(store, trie_cache.clone(), shard_uid, false, None); + let trie_caching_storage = TrieCachingStorage::new( + store, + trie_cache.clone(), + shard_uid, + false, + None, + state_reader, + ); let mut accounting_cache = TrieAccountingCache::new(None); let value = &values[0]; diff --git a/integration-tests/src/tests/client/state_snapshot.rs b/integration-tests/src/tests/client/state_snapshot.rs index 6d1e6b7f85b..085759e3b0d 100644 --- a/integration-tests/src/tests/client/state_snapshot.rs +++ b/integration-tests/src/tests/client/state_snapshot.rs @@ -10,6 +10,7 @@ use near_primitives::shard_layout::ShardUId; use near_primitives::transaction::SignedTransaction; use near_store::config::StateSnapshotType; use near_store::flat::FlatStorageManager; +use near_store::trie::StateReader; use near_store::{ config::TrieCacheConfig, test_utils::create_test_store, Mode, ShardTries, StateSnapshotConfig, StoreConfig, TrieConfig, @@ -50,10 +51,12 @@ impl StateSnaptshotTestEnv { hot_store_path: hot_store_path.clone(), state_snapshot_subdir: state_snapshot_subdir.clone(), }; + let state_reader = StateReader::new_without_resharding_history(&shard_uids); let shard_tries = ShardTries::new( store.clone(), trie_config, &shard_uids, + state_reader, flat_storage_manager, state_snapshot_config, ); diff --git a/nearcore/src/entity_debug.rs b/nearcore/src/entity_debug.rs index be88e6ab04c..bad95d24def 100644 --- a/nearcore/src/entity_debug.rs +++ b/nearcore/src/entity_debug.rs @@ -35,10 +35,10 @@ use near_store::flat::delta::KeyForFlatStateDelta; use near_store::flat::store_helper::encode_flat_state_db_key; use near_store::flat::{FlatStateChanges, FlatStateDeltaMetadata, FlatStorageStatus}; use near_store::{ - DBCol, NibbleSlice, RawTrieNode, RawTrieNodeWithSize, ShardUId, Store, TrieCachingStorage, - CHUNK_TAIL_KEY, COLD_HEAD_KEY, FINAL_HEAD_KEY, FORK_TAIL_KEY, GENESIS_JSON_HASH_KEY, - GENESIS_STATE_ROOTS_KEY, HEADER_HEAD_KEY, HEAD_KEY, LARGEST_TARGET_HEIGHT_KEY, - LATEST_KNOWN_KEY, STATE_SNAPSHOT_KEY, STATE_SYNC_DUMP_KEY, TAIL_KEY, + get_key_from_shard_uid_db_prefix_and_hash, DBCol, NibbleSlice, RawTrieNode, + RawTrieNodeWithSize, ShardUId, Store, CHUNK_TAIL_KEY, COLD_HEAD_KEY, FINAL_HEAD_KEY, + FORK_TAIL_KEY, GENESIS_JSON_HASH_KEY, GENESIS_STATE_ROOTS_KEY, HEADER_HEAD_KEY, HEAD_KEY, + LARGEST_TARGET_HEIGHT_KEY, LATEST_KNOWN_KEY, STATE_SNAPSHOT_KEY, STATE_SYNC_DUMP_KEY, TAIL_KEY, }; use serde::Serialize; use std::collections::{HashMap, HashSet}; @@ -245,13 +245,11 @@ impl EntityDebugHandlerImpl { Ok(serialize_entity(&ExecutionOutcomeView::from(outcome.outcome))) } EntityQuery::RawTrieNodeByHash { trie_node_hash, shard_uid } => { + let state_reader = self.runtime.get_tries().get_state_reader(); let node = store .get_ser::( DBCol::State, - &TrieCachingStorage::get_key_from_shard_uid_and_hash( - shard_uid, - &trie_node_hash, - ), + &state_reader.get_key_from_shard_uid_and_hash(shard_uid, &trie_node_hash), )? .ok_or_else(|| anyhow!("Trie node not found"))?; Ok(serialize_raw_trie_node(node)) @@ -267,25 +265,22 @@ impl EntityDebugHandlerImpl { .shard_uids() .nth(chunk.shard_id() as usize) .ok_or_else(|| anyhow!("Shard {} not found", chunk.shard_id()))?; + let state_reader = self.runtime.get_tries().get_state_reader(); let node = store .get_ser::( DBCol::State, - &TrieCachingStorage::get_key_from_shard_uid_and_hash( - shard_uid, - &chunk.prev_state_root(), - ), + &state_reader + .get_key_from_shard_uid_and_hash(shard_uid, &chunk.prev_state_root()), )? .ok_or_else(|| anyhow!("State root not found"))?; Ok(serialize_raw_trie_node(node)) } EntityQuery::RawTrieValueByHash { trie_value_hash, shard_uid } => { + let state_reader = self.runtime.get_tries().get_state_reader(); let value = store .get( DBCol::State, - &TrieCachingStorage::get_key_from_shard_uid_and_hash( - shard_uid, - &trie_value_hash, - ), + &state_reader.get_key_from_shard_uid_and_hash(shard_uid, &trie_value_hash), )? .ok_or_else(|| anyhow!("Trie value not found"))?; Ok(serialize_entity(&hex::encode(value.as_slice()))) @@ -457,10 +452,7 @@ impl EntityDebugHandlerImpl { .get( DBCol::State, // TODO(reshardingV3) Use shardUId mapping? - &TrieCachingStorage::get_key_from_shard_uid_db_prefix_and_hash( - shard_uid.into(), - &value.hash, - ), + &get_key_from_shard_uid_db_prefix_and_hash(shard_uid.into(), &value.hash), )? .ok_or_else(|| anyhow!("ValueRef could not be dereferenced"))? .to_vec(), diff --git a/nearcore/src/metrics.rs b/nearcore/src/metrics.rs index 3ace5e8d90b..e3426c43ae0 100644 --- a/nearcore/src/metrics.rs +++ b/nearcore/src/metrics.rs @@ -8,7 +8,8 @@ use near_o11y::metrics::{ try_create_int_gauge, try_create_int_gauge_vec, HistogramVec, IntCounterVec, IntGauge, IntGaugeVec, }; -use near_primitives::{shard_layout::ShardLayout, state_record::StateRecord, trie_key}; +use near_primitives::{state_record::StateRecord, trie_key}; +use near_store::trie::StateReader; use near_store::{ShardUId, Store, Trie, TrieDBStorage}; use std::sync::Arc; use std::sync::LazyLock; @@ -125,14 +126,9 @@ fn export_postponed_receipt_count(near_config: &NearConfig, store: &Store) -> an POSTPONED_RECEIPTS_COUNT.with_label_values(&[&shard_id.to_string()]).set(0); continue; } + let shard_uid = ShardUId::from_shard_id_and_layout(shard_id, &shard_layout); - let count = get_postponed_receipt_count_for_shard( - shard_id, - &shard_layout, - &chain_store, - &block, - store, - ); + let count = get_postponed_receipt_count_for_shard(shard_uid, &chain_store, &block, store); let count = match count { Ok(count) => count, Err(err) => { @@ -147,16 +143,15 @@ fn export_postponed_receipt_count(near_config: &NearConfig, store: &Store) -> an } fn get_postponed_receipt_count_for_shard( - shard_id: u64, - shard_layout: &ShardLayout, + shard_uid: ShardUId, chain_store: &ChainStore, block: &Block, store: &Store, ) -> Result { - let shard_uid = ShardUId::from_shard_id_and_layout(shard_id, shard_layout); let chunk_extra = chain_store.get_chunk_extra(block.hash(), &shard_uid)?; let state_root = chunk_extra.state_root(); - let storage = TrieDBStorage::new(store.clone(), shard_uid); + let state_reader = StateReader::new_without_resharding_history(&[shard_uid]); + let storage = TrieDBStorage::new(store.clone(), shard_uid, state_reader); let storage = Arc::new(storage); let flat_storage_chunk_view = None; let trie = Trie::new(storage, *state_root, flat_storage_chunk_view); diff --git a/runtime/runtime-params-estimator/src/estimator_context.rs b/runtime/runtime-params-estimator/src/estimator_context.rs index 071527c26eb..201e8fb6955 100644 --- a/runtime/runtime-params-estimator/src/estimator_context.rs +++ b/runtime/runtime-params-estimator/src/estimator_context.rs @@ -19,6 +19,7 @@ use near_store::flat::{ store_helper, BlockInfo, FlatStateChanges, FlatStateDelta, FlatStateDeltaMetadata, FlatStorage, FlatStorageManager, FlatStorageReadyStatus, FlatStorageStatus, }; +use near_store::trie::StateReader; use near_store::{ShardTries, ShardUId, StateSnapshotConfig, TrieUpdate}; use near_store::{TrieCache, TrieCachingStorage, TrieConfig}; use near_vm_runner::logic::LimitConfig; @@ -92,16 +93,19 @@ impl<'c> EstimatorContext<'c> { let flat_storage = flat_storage_manager.get_flat_storage_for_shard(shard_uid).unwrap(); self.generate_deltas(&flat_storage); + let shard_uids = vec![shard_uid]; // Create ShardTries with relevant settings adjusted for estimator. let mut trie_config = near_store::TrieConfig::default(); trie_config.enable_receipt_prefetching = true; if self.config.memtrie { - trie_config.load_mem_tries_for_shards = vec![shard_uid]; + trie_config.load_mem_tries_for_shards = shard_uids.clone(); } + let state_reader = StateReader::new_without_resharding_history(&shard_uids); let tries = ShardTries::new( store, trie_config, - &[shard_uid], + &shard_uids, + state_reader, flat_storage_manager, StateSnapshotConfig::default(), ); @@ -310,6 +314,7 @@ impl Testbed<'_> { pub(crate) fn trie_caching_storage(&mut self) -> TrieCachingStorage { let store = self.tries.get_store(); + let state_reader = self.tries.get_state_reader(); let is_view = false; let prefetcher = None; let caching_storage = TrieCachingStorage::new( @@ -318,6 +323,7 @@ impl Testbed<'_> { ShardUId::single_shard(), is_view, prefetcher, + state_reader, ); caching_storage } diff --git a/runtime/runtime/src/prefetch.rs b/runtime/runtime/src/prefetch.rs index e75e51f83ce..301978cffdc 100644 --- a/runtime/runtime/src/prefetch.rs +++ b/runtime/runtime/src/prefetch.rs @@ -387,6 +387,7 @@ mod tests { use super::TriePrefetcher; use near_primitives::{trie_key::TrieKey, types::AccountId}; use near_store::test_utils::{create_test_store, test_populate_trie}; + use near_store::trie::StateReader; use near_store::{ShardTries, ShardUId, StateSnapshotConfig, Trie, TrieConfig}; use std::str::FromStr; use std::time::{Duration, Instant}; @@ -474,11 +475,13 @@ mod tests { let shard_uids = vec![ShardUId::single_shard()]; let trie_config = TrieConfig { enable_receipt_prefetching: true, ..TrieConfig::default() }; let store = create_test_store(); + let state_reader = StateReader::new_without_resharding_history(&shard_uids); let flat_storage_manager = near_store::flat::FlatStorageManager::new(store.clone()); let tries = ShardTries::new( store, trie_config, &shard_uids, + state_reader, flat_storage_manager, StateSnapshotConfig::default(), ); diff --git a/tools/database/src/analyze_contract_sizes.rs b/tools/database/src/analyze_contract_sizes.rs index 60f8cc85d8d..dd3c9deb4df 100644 --- a/tools/database/src/analyze_contract_sizes.rs +++ b/tools/database/src/analyze_contract_sizes.rs @@ -5,6 +5,7 @@ use near_chain_configs::GenesisValidationMode; use near_epoch_manager::EpochManager; use near_primitives::trie_key::col; use near_primitives::types::AccountId; +use near_store::trie::StateReader; use near_store::{ShardUId, Trie, TrieDBStorage}; use nearcore::{load_config, open_storage}; use std::collections::BTreeMap; @@ -90,16 +91,19 @@ impl AnalyzeContractSizesCommand { EpochManager::new_from_genesis_config(store.clone(), &near_config.genesis.config) .unwrap(); let shard_layout = epoch_manager.get_shard_layout(&head.epoch_id).unwrap(); + let shard_uids = shard_layout.shard_uids().collect::>(); + let state_reader = StateReader::new_without_resharding_history(&shard_uids); let mut stats = ContractSizeStats::new(self.topn); - for shard_uid in shard_layout.shard_uids() { + for shard_uid in shard_uids { println!("Analyzing chunk with uid: {}", shard_uid); let chunk_extra = chain_store.get_chunk_extra(&head.last_block_hash, &shard_uid).unwrap(); let state_root = chunk_extra.state_root(); - let trie_storage = Arc::new(TrieDBStorage::new(store.clone(), shard_uid)); + let trie_storage = + Arc::new(TrieDBStorage::new(store.clone(), shard_uid, state_reader.clone())); let trie = Trie::new(trie_storage, *state_root, None); let mut iterator = trie.disk_iter().unwrap(); diff --git a/tools/database/src/analyze_delayed_receipt.rs b/tools/database/src/analyze_delayed_receipt.rs index 13942a8d0b1..8b330da92d5 100644 --- a/tools/database/src/analyze_delayed_receipt.rs +++ b/tools/database/src/analyze_delayed_receipt.rs @@ -1,5 +1,6 @@ use clap::Parser; use near_store::flat::FlatStorageManager; +use near_store::trie::StateReader; use near_store::{get_delayed_receipt_indices, ShardTries, StateSnapshotConfig, TrieConfig}; use std::collections::HashMap; use std::path::PathBuf; @@ -57,10 +58,12 @@ impl AnalyzeDelayedReceiptCommand { let tip = chain_store.head().unwrap(); let shard_layout = epoch_manager.get_shard_layout(&tip.epoch_id).unwrap(); let shard_uids = shard_layout.shard_uids().collect::>(); + let state_reader = StateReader::new_without_resharding_history(&shard_uids); let shard_tries = ShardTries::new( store.clone(), TrieConfig::default(), &shard_uids, + state_reader, FlatStorageManager::new(store), StateSnapshotConfig::default(), ); diff --git a/tools/database/src/state_perf.rs b/tools/database/src/state_perf.rs index 9e6af54ccef..a6771976ab3 100644 --- a/tools/database/src/state_perf.rs +++ b/tools/database/src/state_perf.rs @@ -1,5 +1,6 @@ use clap::Parser; use indicatif::{ProgressBar, ProgressIterator}; +use near_store::trie::StateReader; use std::collections::BTreeMap; use std::fmt::{Display, Write}; use std::path::Path; @@ -34,13 +35,17 @@ impl StatePerfCommand { pub(crate) fn run(&self, home: &Path) -> anyhow::Result<()> { let rocksdb = Arc::new(open_rocksdb(home, near_store::Mode::ReadOnly)?); let store = near_store::NodeStorage::new(rocksdb).get_hot_store(); + // TODO(reshardingV3) Need to read all shard UIDs from store, + // or do not require shard UIDs in StateReader constructor at all. + let state_reader = StateReader::new_without_resharding_history(&[]); eprintln!("Start State perf test"); let mut perf_context = PerfContext::new(); let total_samples = self.warmup_samples + self.samples; for (sample_i, (shard_uid, value_ref)) in generate_state_requests(store.clone(), total_samples).into_iter().enumerate().progress() { - let trie_storage = near_store::TrieDBStorage::new(store.clone(), shard_uid); + let trie_storage = + near_store::TrieDBStorage::new(store.clone(), shard_uid, state_reader.clone()); let include_sample = sample_i >= self.warmup_samples; if include_sample { perf_context.reset(); diff --git a/tools/fork-network/src/cli.rs b/tools/fork-network/src/cli.rs index 731b7fe7a6b..4a5ba5c6a41 100644 --- a/tools/fork-network/src/cli.rs +++ b/tools/fork-network/src/cli.rs @@ -27,6 +27,7 @@ use near_primitives::types::{ use near_primitives::version::{ProtocolVersion, PROTOCOL_VERSION}; use near_store::db::RocksDB; use near_store::flat::{store_helper, BlockInfo, FlatStorageManager, FlatStorageStatus}; +use near_store::trie::StateReader; use near_store::{ checkpoint_hot_storage_and_cleanup_columns, DBCol, Store, TrieDBStorage, TrieStorage, FINAL_HEAD_KEY, @@ -504,8 +505,9 @@ impl ForkNetworkCommand { let default_key = near_mirror::key_mapping::default_extra_key(None).public_key(); // Keeps track of accounts that have a full access key. let mut has_full_key = HashSet::new(); + let state_reader = StateReader::new_without_resharding_history(&[shard_uid]); // Lets us lookup large values in the `State` columns. - let trie_storage = TrieDBStorage::new(store.clone(), shard_uid); + let trie_storage = TrieDBStorage::new(store.clone(), shard_uid, state_reader); // Iterate over the whole flat storage and do the necessary changes to have access to all accounts. let mut index_delayed_receipt = 0; diff --git a/tools/state-viewer/src/commands.rs b/tools/state-viewer/src/commands.rs index 6c61cbdc109..2934952bd3c 100644 --- a/tools/state-viewer/src/commands.rs +++ b/tools/state-viewer/src/commands.rs @@ -42,6 +42,7 @@ use near_primitives::version::PROTOCOL_VERSION; use near_primitives_core::types::{Balance, EpochHeight}; use near_store::flat::FlatStorageChunkView; use near_store::flat::FlatStorageManager; +use near_store::trie::StateReader; use near_store::TrieStorage; use near_store::{DBCol, Store, Trie, TrieCache, TrieCachingStorage, TrieConfig, TrieDBStorage}; use nearcore::NightshadeRuntimeExt; @@ -1097,7 +1098,9 @@ fn get_trie(store: Store, hash: CryptoHash, shard_id: u32, shard_version: u32) - let shard_uid = ShardUId { version: shard_version, shard_id }; let trie_config: TrieConfig = Default::default(); let shard_cache = TrieCache::new(&trie_config, shard_uid, true); - let trie_storage = TrieCachingStorage::new(store, shard_cache, shard_uid, true, None); + let state_reader = StateReader::new_without_resharding_history(&[shard_uid]); + let trie_storage = + TrieCachingStorage::new(store, shard_cache, shard_uid, true, None, state_reader); Trie::new(Arc::new(trie_storage), hash, None) } @@ -1174,8 +1177,9 @@ pub(crate) fn contract_accounts( shard_id as u64, &ShardLayout::get_simple_nightshade_layout(), ); + let state_reader = StateReader::new_without_resharding_history(&[shard_uid]); // Use simple non-caching storage, we don't expect many duplicate lookups while iterating. - let storage = TrieDBStorage::new(store.clone(), shard_uid); + let storage = TrieDBStorage::new(store.clone(), shard_uid, state_reader); // We don't need flat state to traverse all accounts. let flat_storage_chunk_view = None; Trie::new(Arc::new(storage), state_root, flat_storage_chunk_view) @@ -1265,7 +1269,8 @@ fn print_state_stats_for_shard_uid( shard_uid: ShardUId, ) { flat_storage_manager.create_flat_storage_for_shard(shard_uid).unwrap(); - let trie_storage = TrieDBStorage::new(store.clone(), shard_uid); + let state_reader = StateReader::new_without_resharding_history(&[shard_uid]); + let trie_storage = TrieDBStorage::new(store.clone(), shard_uid, state_reader); let chunk_view = flat_storage_manager.chunk_view(shard_uid, block_hash).unwrap(); let mut state_stats = StateStats::default(); diff --git a/tools/state-viewer/src/trie_iteration_benchmark.rs b/tools/state-viewer/src/trie_iteration_benchmark.rs index ea6c211fdb3..53793daf142 100644 --- a/tools/state-viewer/src/trie_iteration_benchmark.rs +++ b/tools/state-viewer/src/trie_iteration_benchmark.rs @@ -8,6 +8,7 @@ use near_primitives::trie_key::trie_key_parsers::{ parse_account_id_from_access_key_key, parse_account_id_from_trie_key_with_separator, }; use near_primitives_core::types::ShardId; +use near_store::trie::StateReader; use near_store::{ShardUId, Store, Trie, TrieDBStorage}; use nearcore::NearConfig; use std::cell::RefCell; @@ -148,7 +149,8 @@ impl TrieIterationBenchmarkCmd { // corresponds to the current epoch id. In practice shouldn't // matter as the shard layout doesn't change. let state_root = chunk_header.prev_state_root(); - let storage = TrieDBStorage::new(store.clone(), shard_uid); + let state_reader = StateReader::new_without_resharding_history(&[shard_uid]); + let storage = TrieDBStorage::new(store.clone(), shard_uid, state_reader); let flat_storage_chunk_view = None; Trie::new(Arc::new(storage), state_root, flat_storage_chunk_view) }