From a85db92e8d160eac3bd41615d25bc9fdee6cd5bc Mon Sep 17 00:00:00 2001 From: Aleksandr Logunov Date: Tue, 9 Aug 2022 19:44:02 +0400 Subject: [PATCH] feat: simple flat storage (#7345) First step towards flat storage: https://github.com/near/nearcore/issues/7327 Continue work here because https://github.com/near/nearcore/pull/7295 was broken. Here we start building it under protocol feature without adding it to nightly. The goal is to support it for localnet with one node which never sends skip approvals, so later we could extend it to more nodes and eventually implement a migration to mainnet. Now, everyone should be able to build neard with this feature, set `max_block_production_delay` to something big and run localnet. I currently couldn't make all tests work with enabled feature, e.g. `test_care_about_shard`, but I open it for review because I don't think it affects the idea. ## Testing * Check that encoding and decoding value reference gives expected results. `ValueRef` can be reused in other places in code * Check that flat state is used in regular trie handler and not used in view trie handler * Check that after block processing, getting value from regular trie and view trie gives the same result --- chain/chain/Cargo.toml | 1 + chain/chain/src/store.rs | 4 ++ core/primitives/src/lib.rs | 1 + core/primitives/src/state.rs | 49 ++++++++++++++++++++ core/primitives/src/state_record.rs | 6 +++ core/store/Cargo.toml | 2 + core/store/src/columns.rs | 7 +++ core/store/src/flat_state.rs | 64 ++++++++++++++++++++++++++ core/store/src/lib.rs | 1 + core/store/src/trie/mod.rs | 40 ++++++++++------- core/store/src/trie/shard_tries.rs | 70 ++++++++++++++++++++++++++--- core/store/src/trie/trie_tests.rs | 2 +- core/store/src/trie/update.rs | 25 +++-------- nearcore/Cargo.toml | 2 + nearcore/src/runtime/mod.rs | 65 ++++++++++++++++++++++++++- neard/Cargo.toml | 2 + runtime/runtime/Cargo.toml | 1 + runtime/runtime/src/genesis.rs | 18 +++++++- 18 files changed, 316 insertions(+), 44 deletions(-) create mode 100644 core/primitives/src/state.rs create mode 100644 core/store/src/flat_state.rs diff --git a/chain/chain/Cargo.toml b/chain/chain/Cargo.toml index a6a758909f5..5d2d5563f7b 100644 --- a/chain/chain/Cargo.toml +++ b/chain/chain/Cargo.toml @@ -53,6 +53,7 @@ protocol_feature_chunk_only_producers = [ "near-chain-configs/protocol_feature_chunk_only_producers", "near-primitives/protocol_feature_chunk_only_producers", ] +protocol_feature_flat_state = ["near-store/protocol_feature_flat_state"] nightly = [ "nightly_protocol", diff --git a/chain/chain/src/store.rs b/chain/chain/src/store.rs index ee763e1ba13..dc03847a0d6 100644 --- a/chain/chain/src/store.rs +++ b/chain/chain/src/store.rs @@ -2440,6 +2440,10 @@ impl<'a> ChainStoreUpdate<'a> { | DBCol::CachedContractCode => { unreachable!(); } + #[cfg(feature = "protocol_feature_flat_state")] + DBCol::FlatState => { + unreachable!(); + } } self.inc_gc(col); self.merge(store_update); diff --git a/core/primitives/src/lib.rs b/core/primitives/src/lib.rs index efb2b599c70..d1e08910dcc 100644 --- a/core/primitives/src/lib.rs +++ b/core/primitives/src/lib.rs @@ -21,6 +21,7 @@ pub mod runtime; pub mod sandbox_state_patch; pub mod shard_layout; pub mod sharding; +pub mod state; pub mod state_part; pub mod state_record; pub mod syncing; diff --git a/core/primitives/src/state.rs b/core/primitives/src/state.rs new file mode 100644 index 00000000000..42bbc3030bd --- /dev/null +++ b/core/primitives/src/state.rs @@ -0,0 +1,49 @@ +use byteorder::{LittleEndian, ReadBytesExt}; +use near_primitives_core::hash::{hash, CryptoHash}; +use std::io::{Cursor, Read}; + +/// State value reference. Used to charge fees for value length before retrieving the value itself. +pub struct ValueRef { + /// Value length in bytes. + pub length: u32, + /// Unique value hash. + pub hash: CryptoHash, +} + +impl ValueRef { + /// Create serialized value reference by the value. + /// Resulting array stores 4 bytes of length and then 32 bytes of hash. + /// TODO (#7327): consider passing hash here to avoid double computation + pub fn create_serialized(value: &[u8]) -> [u8; 36] { + let mut result = [0u8; 36]; + result[0..4].copy_from_slice(&(value.len() as u32).to_le_bytes()); + result[4..36].copy_from_slice(&hash(value).0); + result + } + + /// Decode value reference from the raw byte array. + /// TODO (#7327): use &[u8; 36] and get rid of Cursor; also check that there are no leftover bytes + pub fn decode(bytes: &[u8]) -> Result { + let mut cursor = Cursor::new(bytes); + let value_length = cursor.read_u32::()?; + let mut arr = [0; 32]; + cursor.read_exact(&mut arr)?; + let value_hash = CryptoHash(arr); + Ok(ValueRef { length: value_length, hash: value_hash }) + } +} + +#[cfg(test)] +mod tests { + use crate::state::ValueRef; + use near_primitives_core::hash::hash; + + #[test] + fn test_encode_decode() { + let value = vec![1, 2, 3]; + let value_ref_ser = ValueRef::create_serialized(&value); + let value_ref = ValueRef::decode(&value_ref_ser).unwrap(); + assert_eq!(value_ref.length, value.len() as u32); + assert_eq!(value_ref.hash, hash(&value)); + } +} diff --git a/core/primitives/src/state_record.rs b/core/primitives/src/state_record.rs index 41540500cad..1d0ba416de3 100644 --- a/core/primitives/src/state_record.rs +++ b/core/primitives/src/state_record.rs @@ -161,5 +161,11 @@ pub fn state_record_to_account_id(state_record: &StateRecord) -> &AccountId { } pub fn is_contract_code_key(key: &[u8]) -> bool { + debug_assert!(!key.is_empty()); key[0] == col::CONTRACT_CODE } + +pub fn is_delayed_receipt_key(key: &[u8]) -> bool { + debug_assert!(!key.is_empty()); + key[0] == col::DELAYED_RECEIPT || key[0] == col::DELAYED_RECEIPT_INDICES +} diff --git a/core/store/Cargo.toml b/core/store/Cargo.toml index 00e5d490678..d79c810e8b7 100644 --- a/core/store/Cargo.toml +++ b/core/store/Cargo.toml @@ -54,6 +54,8 @@ no_cache = [] single_thread_rocksdb = [] # Deactivate RocksDB IO background threads test_features = [] protocol_feature_chunk_only_producers = [] +protocol_feature_flat_state = [] + nightly_protocol = [] nightly = [ "nightly_protocol", diff --git a/core/store/src/columns.rs b/core/store/src/columns.rs index 1cf518c1a6d..98e14d3b8fa 100644 --- a/core/store/src/columns.rs +++ b/core/store/src/columns.rs @@ -234,6 +234,11 @@ pub enum DBCol { /// - *Rows*: BlockShardId (BlockHash || ShardId) - 40 bytes /// - *Column type*: StateChangesForSplitStates StateChangesForSplitStates = 49, + /// State changes made by a chunk, used for splitting states + /// - *Rows*: serialized TrieKey (Vec) + /// - *Column type*: ValueRef + #[cfg(feature = "protocol_feature_flat_state")] + FlatState = 50, } impl DBCol { @@ -406,6 +411,8 @@ impl fmt::Display for DBCol { Self::EpochValidatorInfo => "epoch validator info", Self::HeaderHashesByHeight => "header hashes indexed by their height", Self::StateChangesForSplitStates => "state changes indexed by block hash and shard id", + #[cfg(feature = "protocol_feature_flat_state")] + Self::FlatState => "flat state", }; write!(f, "{}", desc) } diff --git a/core/store/src/flat_state.rs b/core/store/src/flat_state.rs new file mode 100644 index 00000000000..43df74889c6 --- /dev/null +++ b/core/store/src/flat_state.rs @@ -0,0 +1,64 @@ +//! Contains flat state optimization logic. +//! +//! The state of the contract is a key-value map, `Map, Vec>`. +//! In the database, we store this map as a trie, which allows us to construct succinct proofs that a certain key/value +//! belongs to contract's state. Using a trie has a drawback -- reading a single key/value requires traversing the trie +//! from the root, loading many nodes from the database. +//! To optimize this, we want to use flat state: alongside the trie, we store a mapping from keys to value +//! references so that, if you don't need a proof, you can do a db lookup in just two db accesses - one to get value +//! reference, one to get value itself. +/// TODO (#7327): consider inlining small values, so we could use only one db access. + +#[cfg(feature = "protocol_feature_flat_state")] +use crate::DBCol; +use crate::Store; +use near_primitives::errors::StorageError; +use near_primitives::hash::CryptoHash; +use near_primitives::state::ValueRef; + +/// Struct for getting value references from the flat storage. +/// Used to speed up `get` and `get_ref` trie methods. +/// It should store all trie keys for state on top of chain head, except delayed receipt keys, because they are the +/// same for each shard and they are requested only once during applying chunk. +/// TODO (#7327): implement flat state deltas to support forks +/// TODO (#7327): store on top of final head (or earlier) so updates will only go forward +#[derive(Clone)] +pub struct FlatState { + #[allow(dead_code)] + store: Store, +} + +impl FlatState { + #[cfg(feature = "protocol_feature_flat_state")] + pub fn new(store: Store) -> Self { + Self { store } + } + + #[allow(unused_variables)] + fn get_raw_ref(&self, key: &[u8]) -> Result>, StorageError> { + #[cfg(feature = "protocol_feature_flat_state")] + return self + .store + .get(DBCol::FlatState, key) + .map_err(|_| StorageError::StorageInternalError); + #[cfg(not(feature = "protocol_feature_flat_state"))] + unreachable!(); + } + + /// Get value reference using raw trie key and state root. We assume that flat state contains data for this root. + /// To avoid duplication, we don't store values themselves in flat state, they are stored in `DBCol::State`. Also + /// the separation is done so we could charge users for the value length before loading the value. + /// TODO (#7327): support different roots (or block hashes). + pub fn get_ref( + &self, + _root: &CryptoHash, + key: &[u8], + ) -> Result, StorageError> { + match self.get_raw_ref(key)? { + Some(bytes) => { + ValueRef::decode(&bytes).map(Some).map_err(|_| StorageError::StorageInternalError) + } + None => Ok(None), + } + } +} diff --git a/core/store/src/lib.rs b/core/store/src/lib.rs index 82c724d3634..37482a81a48 100644 --- a/core/store/src/lib.rs +++ b/core/store/src/lib.rs @@ -39,6 +39,7 @@ pub use crate::trie::{ mod columns; mod config; pub mod db; +pub mod flat_state; mod metrics; pub mod migrations; pub mod test_utils; diff --git a/core/store/src/trie/mod.rs b/core/store/src/trie/mod.rs index 9c70247eb86..8d48bee2283 100644 --- a/core/store/src/trie/mod.rs +++ b/core/store/src/trie/mod.rs @@ -9,8 +9,11 @@ use near_primitives::challenge::PartialState; use near_primitives::contract::ContractCode; use near_primitives::hash::{hash, CryptoHash}; pub use near_primitives::shard_layout::ShardUId; +use near_primitives::state::ValueRef; +use near_primitives::state_record::is_delayed_receipt_key; use near_primitives::types::{StateRoot, StateRootNode}; +use crate::flat_state::FlatState; use crate::trie::insert_delete::NodesStorage; use crate::trie::iterator::TrieIterator; use crate::trie::nibble_slice::NibbleSlice; @@ -403,7 +406,8 @@ impl RawTrieNodeWithSize { } pub struct Trie { - pub(crate) storage: Box, + pub storage: Box, + pub flat_state: Option, } /// Stores reference count change for some key-value pair in DB. @@ -468,8 +472,8 @@ pub struct ApplyStatePartResult { impl Trie { pub const EMPTY_ROOT: StateRoot = StateRoot::new(); - pub fn new(store: Box) -> Self { - Trie { storage: store } + pub fn new(storage: Box, flat_state: Option) -> Self { + Trie { storage, flat_state } } pub fn recording_reads(&self) -> Self { @@ -480,7 +484,7 @@ impl Trie { shard_uid: storage.shard_uid, recorded: RefCell::new(Default::default()), }; - Trie { storage: Box::new(storage) } + Trie { storage: Box::new(storage), flat_state: None } } pub fn recorded_storage(&self) -> Option { @@ -499,6 +503,7 @@ impl Trie { recorded_storage, visited_nodes: Default::default(), }), + flat_state: None, } } @@ -612,7 +617,7 @@ impl Trie { &self, root: &CryptoHash, mut key: NibbleSlice<'_>, - ) -> Result, StorageError> { + ) -> Result, StorageError> { let mut hash = *root; loop { let node = match self.retrieve_raw_node(&hash)? { @@ -622,7 +627,7 @@ impl Trie { match node { RawTrieNode::Leaf(existing_key, value_length, value_hash) => { if NibbleSlice::from_encoded(&existing_key).0 == key { - return Ok(Some((value_length, value_hash))); + return Ok(Some(ValueRef { length: value_length, hash: value_hash })); } else { return Ok(None); } @@ -640,7 +645,10 @@ impl Trie { if key.is_empty() { match value { Some((value_length, value_hash)) => { - return Ok(Some((value_length, value_hash))); + return Ok(Some(ValueRef { + length: value_length, + hash: value_hash, + })); } None => return Ok(None), } @@ -658,18 +666,20 @@ impl Trie { } } - pub fn get_ref( - &self, - root: &CryptoHash, - key: &[u8], - ) -> Result, StorageError> { - let key = NibbleSlice::new(key); - self.lookup(root, key) + pub fn get_ref(&self, root: &CryptoHash, key: &[u8]) -> Result, StorageError> { + let is_delayed = is_delayed_receipt_key(key); + match &self.flat_state { + Some(flat_state) if !is_delayed => flat_state.get_ref(root, &key), + _ => { + let key = NibbleSlice::new(key); + self.lookup(root, key) + } + } } pub fn get(&self, root: &CryptoHash, key: &[u8]) -> Result>, StorageError> { match self.get_ref(root, key)? { - Some((_length, hash)) => { + Some(ValueRef { hash, .. }) => { self.storage.retrieve_raw_bytes(&hash).map(|bytes| Some(bytes.to_vec())) } None => Ok(None), diff --git a/core/store/src/trie/shard_tries.rs b/core/store/src/trie/shard_tries.rs index f592630a022..edbcd58b1b1 100644 --- a/core/store/src/trie/shard_tries.rs +++ b/core/store/src/trie/shard_tries.rs @@ -12,6 +12,13 @@ use near_primitives::types::{ NumShards, RawStateChange, RawStateChangesWithTrieKey, StateChangeCause, StateRoot, }; +#[cfg(feature = "protocol_feature_flat_state")] +use near_primitives::state::ValueRef; +#[cfg(feature = "protocol_feature_flat_state")] +use near_primitives::state_record::is_delayed_receipt_key; + +#[cfg(feature = "protocol_feature_flat_state")] +use crate::flat_state::FlatState; use crate::trie::trie_storage::{TrieCache, TrieCachingStorage}; use crate::trie::{TrieRefcountChange, POISONED_LOCK_ERR}; use crate::{DBCol, DBOp, DBTransaction}; @@ -92,7 +99,13 @@ impl ShardTries { TrieUpdate::new(Rc::new(self.get_view_trie_for_shard(shard_uid)), state_root) } - fn get_trie_for_shard_internal(&self, shard_uid: ShardUId, is_view: bool) -> Trie { + #[allow(unused_variables)] + fn get_trie_for_shard_internal( + &self, + shard_uid: ShardUId, + is_view: bool, + use_flat_state: bool, + ) -> Trie { let caches_to_use = if is_view { &self.0.view_caches } else { &self.0.caches }; let cache = { let mut caches = caches_to_use.write().expect(POISONED_LOCK_ERR); @@ -101,16 +114,30 @@ impl ShardTries { .or_insert_with(|| self.0.trie_cache_factory.create_cache(&shard_uid)) .clone() }; - let store = Box::new(TrieCachingStorage::new(self.0.store.clone(), cache, shard_uid)); - Trie::new(store) + let storage = Box::new(TrieCachingStorage::new(self.0.store.clone(), cache, shard_uid)); + let flat_state = { + #[cfg(feature = "protocol_feature_flat_state")] + if use_flat_state { + Some(FlatState::new(self.0.store.clone())) + } else { + None + } + #[cfg(not(feature = "protocol_feature_flat_state"))] + None + }; + Trie::new(storage, flat_state) } pub fn get_trie_for_shard(&self, shard_uid: ShardUId) -> Trie { - self.get_trie_for_shard_internal(shard_uid, false) + self.get_trie_for_shard_internal(shard_uid, false, false) + } + + pub fn get_trie_with_flat_state_for_shard(&self, shard_uid: ShardUId) -> Trie { + self.get_trie_for_shard_internal(shard_uid, false, true) } pub fn get_view_trie_for_shard(&self, shard_uid: ShardUId) -> Trie { - self.get_trie_for_shard_internal(shard_uid, true) + self.get_trie_for_shard_internal(shard_uid, true, false) } pub fn get_store(&self) -> Store { @@ -234,6 +261,36 @@ impl ShardTries { ) -> (StoreUpdate, StateRoot) { self.apply_all_inner(trie_changes, shard_uid, true) } + + // TODO(#7327): consider uniting with `apply_all` + #[cfg(feature = "protocol_feature_flat_state")] + pub fn apply_changes_to_flat_state( + &self, + changes: &[RawStateChangesWithTrieKey], + store_update: &mut StoreUpdate, + ) { + for change in changes.iter() { + let key = change.trie_key.to_vec(); + if is_delayed_receipt_key(&key) { + continue; + } + + // `RawStateChangesWithTrieKey` stores all sequential changes for a key within a chunk, so it is sufficient + // to take only the last change. + let last_change = &change + .changes + .last() + .expect("Committed entry should have at least one change") + .data; + match last_change { + Some(value) => { + let value_ref_ser = ValueRef::create_serialized(value); + store_update.set(DBCol::FlatState, &key, &value_ref_ser) + } + None => store_update.delete(DBCol::FlatState, &key), + } + } + } } pub struct WrappedTrieChanges { @@ -267,6 +324,9 @@ impl WrappedTrieChanges { /// /// NOTE: the changes are drained from `self`. pub fn state_changes_into(&mut self, store_update: &mut StoreUpdate) { + #[cfg(feature = "protocol_feature_flat_state")] + self.tries.apply_changes_to_flat_state(&self.state_changes, store_update); + for change_with_trie_key in self.state_changes.drain(..) { assert!( !change_with_trie_key.changes.iter().any(|RawStateChange { cause, .. }| matches!( diff --git a/core/store/src/trie/trie_tests.rs b/core/store/src/trie/trie_tests.rs index 4767cd54464..1b071eb4e6b 100644 --- a/core/store/src/trie/trie_tests.rs +++ b/core/store/src/trie/trie_tests.rs @@ -80,7 +80,7 @@ where print!("Test touches {} nodes, expected result {:?}...", size, expected); for i in 0..(size + 1) { let storage = IncompletePartialStorage::new(storage.clone(), i); - let trie = Trie { storage: Box::new(storage) }; + let trie = Trie { storage: Box::new(storage), flat_state: None }; let expected_result = if i < size { Err(&StorageError::TrieNodeMissing) } else { Ok(&expected) }; assert_eq!(test(Rc::new(trie)).as_ref(), expected_result); diff --git a/core/store/src/trie/update.rs b/core/store/src/trie/update.rs index e79334dae03..446eef90062 100644 --- a/core/store/src/trie/update.rs +++ b/core/store/src/trie/update.rs @@ -10,6 +10,7 @@ use crate::trie::TrieChanges; use crate::StorageError; use super::{Trie, TrieIterator}; +use near_primitives::state::ValueRef; use near_primitives::trie_key::TrieKey; use std::rc::Rc; @@ -23,6 +24,7 @@ pub struct TrieKeyValueUpdate { pub type TrieUpdates = BTreeMap, TrieKeyValueUpdate>; /// Provides a way to access Storage and record changes with future commit. +/// TODO (#7327): rename to StateUpdate pub struct TrieUpdate { pub trie: Rc, root: CryptoHash, @@ -84,8 +86,11 @@ impl TrieUpdate { return Ok(data.as_ref().map(TrieUpdateValuePtr::MemoryRef)); } } + self.trie.get_ref(&self.root, &key).map(|option| { - option.map(|(length, hash)| TrieUpdateValuePtr::HashAndSize(&self.trie, length, hash)) + option.map(|ValueRef { length, hash }| { + TrieUpdateValuePtr::HashAndSize(&self.trie, length, hash) + }) }) } @@ -136,24 +141,6 @@ impl TrieUpdate { Ok((trie_changes, state_changes)) } - pub fn finalize_genesis(self) -> Result { - assert!(self.prospective.is_empty(), "Finalize cannot be called with uncommitted changes."); - let TrieUpdate { trie, root, committed, .. } = self; - let trie_changes = trie.update( - &root, - committed.into_iter().map(|(k, changes_with_trie_key)| { - let data = changes_with_trie_key - .changes - .into_iter() - .last() - .expect("Committed entry should have at least one change") - .data; - (k, data) - }), - )?; - Ok(trie_changes) - } - /// Returns Error if the underlying storage fails pub fn iter(&self, key_prefix: &[u8]) -> Result, StorageError> { TrieUpdateIterator::new(self, key_prefix, b"", None) diff --git a/nearcore/Cargo.toml b/nearcore/Cargo.toml index 4ae8c510010..b15b3b7a9ff 100644 --- a/nearcore/Cargo.toml +++ b/nearcore/Cargo.toml @@ -120,6 +120,8 @@ protocol_feature_fix_staking_threshold = [ protocol_feature_fix_contract_loading_cost = [ "near-vm-runner/protocol_feature_fix_contract_loading_cost", ] +protocol_feature_flat_state = ["near-store/protocol_feature_flat_state", "near-chain/protocol_feature_flat_state", "node-runtime/protocol_feature_flat_state"] + nightly = [ "nightly_protocol", "near-primitives/nightly", diff --git a/nearcore/src/runtime/mod.rs b/nearcore/src/runtime/mod.rs index 9d67b070d4d..41491f77d63 100644 --- a/nearcore/src/runtime/mod.rs +++ b/nearcore/src/runtime/mod.rs @@ -716,9 +716,11 @@ impl RuntimeAdapter for NightshadeRuntime { self.tries.clone() } + // TODO (#7327): Make usage of flat state conditional on prev_hash and call `get_trie_for_shard` if this is not the + // case. Current implementation never creates flat state if `protocol_feature_flat_state` is not enabled. fn get_trie_for_shard(&self, shard_id: ShardId, prev_hash: &CryptoHash) -> Result { let shard_uid = self.get_shard_uid_from_prev_hash(shard_id, prev_hash)?; - Ok(self.tries.get_trie_for_shard(shard_uid)) + Ok(self.tries.get_trie_with_flat_state_for_shard(shard_uid)) } fn get_view_trie_for_shard( @@ -2014,7 +2016,7 @@ mod test { use near_logger_utils::init_test_logger; use near_primitives::block::Tip; use near_primitives::challenge::SlashedValidator; - use near_primitives::transaction::{Action, DeleteAccountAction, StakeAction}; + use near_primitives::transaction::{Action, DeleteAccountAction, StakeAction, TransferAction}; use near_primitives::types::{BlockHeightDelta, Nonce, ValidatorId, ValidatorKickoutReason}; use near_primitives::validator_signer::{InMemoryValidatorSigner, ValidatorSigner}; use near_primitives::views::{ @@ -2025,6 +2027,7 @@ mod test { use super::*; + use near_primitives::trie_key::TrieKey; use primitive_types::U256; fn stake( @@ -2908,6 +2911,7 @@ mod test { #[test] fn test_care_about_shard() { + init_test_logger(); let num_nodes = 2; let validators = (0..num_nodes) .map(|i| AccountId::try_from(format!("test{}", i + 1)).unwrap()) @@ -3482,4 +3486,61 @@ mod test { assert_eq!(env.last_proposals.len(), 1); assert_eq!(env.last_proposals[0].stake(), 0); } + + /// Check that flat state is included into trie and is not included into view trie, because we can't apply flat + /// state optimization to view calls. + #[test] + fn test_flat_state_usage() { + let env = TestEnv::new(vec![vec!["test1".parse().unwrap()]], 4, false); + let trie = env.runtime.get_trie_for_shard(0, &env.head.prev_block_hash).unwrap(); + assert_eq!(trie.flat_state.is_some(), cfg!(feature = "protocol_feature_flat_state")); + + let trie = env.runtime.get_view_trie_for_shard(0, &env.head.prev_block_hash).unwrap(); + assert!(trie.flat_state.is_none()); + } + + /// Check that querying trie and flat state gives the same result. + #[test] + fn test_trie_and_flat_state_equality() { + let num_nodes = 2; + let validators = (0..num_nodes) + .map(|i| AccountId::try_from(format!("test{}", i + 1)).unwrap()) + .collect::>(); + let mut env = TestEnv::new(vec![validators.clone()], 4, false); + let signers: Vec<_> = validators + .iter() + .map(|id| InMemorySigner::from_seed(id.clone(), KeyType::ED25519, id.as_ref())) + .collect(); + + let transfer_tx = SignedTransaction::from_actions( + 4, + signers[0].account_id.clone(), + validators[1].clone(), + &signers[0] as &dyn Signer, + vec![Action::Transfer(TransferAction { deposit: 10 })], + // runtime does not validate block history + CryptoHash::default(), + ); + env.step_default(vec![transfer_tx]); + for _ in 1..=5 { + env.step_default(vec![]); + } + + // Extract account in two ways: + // - using state trie, which should use flat state after enabling it in the protocol + // - using view state, which should never use flat state + let head_prev_block_hash = env.head.prev_block_hash; + let state_root = env.state_roots[0]; + let state = env.runtime.get_trie_for_shard(0, &head_prev_block_hash).unwrap(); + let view_state = env.runtime.get_trie_for_shard(0, &head_prev_block_hash).unwrap(); + let trie_key = TrieKey::Account { account_id: validators[1].clone() }; + let key = trie_key.to_vec(); + + let state_value = state.get(&state_root, &key).unwrap().unwrap(); + let account = Account::try_from_slice(&state_value).unwrap(); + assert_eq!(account.amount(), TESTING_INIT_BALANCE - TESTING_INIT_STAKE + 10); + + let view_state_value = view_state.get(&state_root, &key).unwrap().unwrap(); + assert_eq!(state_value, view_state_value); + } } diff --git a/neard/Cargo.toml b/neard/Cargo.toml index 1cb5e7adb5b..5c34ec31809 100644 --- a/neard/Cargo.toml +++ b/neard/Cargo.toml @@ -63,6 +63,8 @@ protocol_feature_chunk_only_producers = [ "near-primitives/protocol_feature_chunk_only_producers", ] protocol_feature_fix_staking_threshold = ["nearcore/protocol_feature_fix_staking_threshold"] +protocol_feature_flat_state = ["nearcore/protocol_feature_flat_state"] + nightly = [ "nightly_protocol", "nearcore/nightly" diff --git a/runtime/runtime/Cargo.toml b/runtime/runtime/Cargo.toml index 97577c99226..c73d7159e0a 100644 --- a/runtime/runtime/Cargo.toml +++ b/runtime/runtime/Cargo.toml @@ -41,6 +41,7 @@ protocol_feature_chunk_only_producers = [ "near-store/protocol_feature_chunk_only_producers", "near-chain-configs/protocol_feature_chunk_only_producers", ] +protocol_feature_flat_state = ["near-store/protocol_feature_flat_state"] no_cpu_compatibility_checks = ["near-vm-runner/no_cpu_compatibility_checks"] no_cache = [ diff --git a/runtime/runtime/src/genesis.rs b/runtime/runtime/src/genesis.rs index 34d5baf89e6..0835ae02ab5 100644 --- a/runtime/runtime/src/genesis.rs +++ b/runtime/runtime/src/genesis.rs @@ -92,9 +92,23 @@ impl GenesisStateApplier { shard_uid: ShardUId, ) { state_update.commit(StateChangeCause::InitialState); - let trie_changes = state_update.finalize_genesis().expect("Genesis state update failed"); - let (store_update, new_state_root) = tries.apply_all(&trie_changes, shard_uid); + #[cfg(feature = "protocol_feature_flat_state")] + let (store_update, new_state_root) = { + let (trie_changes, state_changes) = + state_update.finalize().expect("Genesis state update failed"); + let (mut store_update, new_state_root) = tries.apply_all(&trie_changes, shard_uid); + tries.apply_changes_to_flat_state(&state_changes, &mut store_update); + (store_update, new_state_root) + }; + + #[cfg(not(feature = "protocol_feature_flat_state"))] + let (store_update, new_state_root) = { + let (trie_changes, _) = state_update.finalize().expect("Genesis state update failed"); + let (store_update, new_state_root) = tries.apply_all(&trie_changes, shard_uid); + (store_update, new_state_root) + }; + store_update.commit().expect("Store update failed on genesis initialization"); *current_state_root = new_state_root; }