From 455ae9432cf5705689cd980400f764980a1bf56b Mon Sep 17 00:00:00 2001 From: Simonas Kazlauskas Date: Thu, 25 Jan 2024 13:49:15 +0200 Subject: [PATCH 1/6] store: add method to check key presence in trie Outstanding question: do we need or want to consult flat storage when checking for presence? --- core/store/src/trie/mod.rs | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/core/store/src/trie/mod.rs b/core/store/src/trie/mod.rs index 03c7216fee8..d6bb3006e99 100644 --- a/core/store/src/trie/mod.rs +++ b/core/store/src/trie/mod.rs @@ -1360,6 +1360,21 @@ impl Trie { Ok(bytes.to_vec()) } + /// Check if the column contains a value with the given `key`. + /// + /// This method is guaranteed to not inspect the value stored for this key, which would + /// otherwise have potential gas cost implications. + pub fn contains_key(&self, key: &[u8]) -> Result { + let charge_gas_for_trie_node_access = self.charge_gas_for_trie_node_access; + if self.memtries.is_some() { + Ok(self.lookup_from_memory(key, charge_gas_for_trie_node_access, |_| ())?.is_some()) + } else { + Ok(self + .lookup_from_state_column(NibbleSlice::new(key), charge_gas_for_trie_node_access)? + .is_some()) + } + } + /// Retrieves an `OptimizedValueRef`` for the given key. See `OptimizedValueRef`. /// /// `mode`: whether we will try to perform the lookup through flat storage or trie. From 86598743507610776c05e28b902194d924a83acf Mon Sep 17 00:00:00 2001 From: Simonas Kazlauskas Date: Thu, 25 Jan 2024 14:37:24 +0200 Subject: [PATCH 2/6] store: thread through contains_key --- core/store/src/trie/mod.rs | 9 +++++++++ core/store/src/trie/update.rs | 16 ++++++++++++++++ 2 files changed, 25 insertions(+) diff --git a/core/store/src/trie/mod.rs b/core/store/src/trie/mod.rs index d6bb3006e99..7f9ad859d7d 100644 --- a/core/store/src/trie/mod.rs +++ b/core/store/src/trie/mod.rs @@ -365,6 +365,11 @@ pub trait TrieAccess { /// root are already known by the object rather than being passed as /// argument. fn get(&self, key: &TrieKey) -> Result>, StorageError>; + + /// Check if the key is present. + /// + /// Equivalent to `Self::get(k)?.is_some()`, but avoids reading out the value. + fn contains_key(&self, key: &TrieKey) -> Result; } /// Stores reference count addition for some key-value pair in DB. @@ -1508,6 +1513,10 @@ impl TrieAccess for Trie { fn get(&self, key: &TrieKey) -> Result>, StorageError> { Trie::get(self, &key.to_vec()) } + + fn contains_key(&self, key: &TrieKey) -> Result { + Trie::contains_key(&self, &key.to_vec()) + } } /// Counts trie nodes reads during tx/receipt execution for proper storage costs charging. diff --git a/core/store/src/trie/update.rs b/core/store/src/trie/update.rs index b3d6fd46a54..eba79d8d589 100644 --- a/core/store/src/trie/update.rs +++ b/core/store/src/trie/update.rs @@ -80,6 +80,18 @@ impl TrieUpdate { Ok(result) } + pub fn contains_key(&self, key: &TrieKey) -> Result { + let key = key.to_vec(); + if self.prospective.contains_key(&key) { + return Ok(true); + } else if let Some(changes_with_trie_key) = self.committed.get(&key) { + if let Some(RawStateChange { data, .. }) = changes_with_trie_key.changes.last() { + return Ok(data.is_some()); + } + } + self.trie.contains_key(&key) + } + pub fn get(&self, key: &TrieKey) -> Result>, StorageError> { let key = key.to_vec(); if let Some(key_value) = self.prospective.get(&key) { @@ -163,6 +175,10 @@ impl crate::TrieAccess for TrieUpdate { fn get(&self, key: &TrieKey) -> Result>, StorageError> { TrieUpdate::get(self, key) } + + fn contains_key(&self, key: &TrieKey) -> Result { + TrieUpdate::contains_key(&self, key) + } } #[cfg(test)] From 76beea56bfc47ce13d45adc8005fafae988848aa Mon Sep 17 00:00:00 2001 From: Simonas Kazlauskas Date: Thu, 25 Jan 2024 14:43:26 +0200 Subject: [PATCH 3/6] perf: add `has_received_data` and use it `get_received_data` has a property of returning the data read out from the storage which can result in the data being read out multiple times only for most of these reads to be immediately discarded. The newly added functions enable explicitly checking for presence of a key without reading out the corresponding value. --- core/store/src/genesis/state_applier.rs | 5 ++--- core/store/src/lib.rs | 8 ++++++++ runtime/runtime/src/lib.rs | 8 ++++---- 3 files changed, 14 insertions(+), 7 deletions(-) diff --git a/core/store/src/genesis/state_applier.rs b/core/store/src/genesis/state_applier.rs index ead0b6b0d20..a11202759ae 100644 --- a/core/store/src/genesis/state_applier.rs +++ b/core/store/src/genesis/state_applier.rs @@ -1,6 +1,6 @@ use crate::flat::FlatStateChanges; use crate::{ - get_account, get_received_data, set, set_access_key, set_account, set_code, + get_account, has_received_data, set, set_access_key, set_account, set_code, set_delayed_receipt, set_postponed_receipt, set_received_data, ShardTries, TrieUpdate, }; @@ -277,9 +277,8 @@ impl GenesisStateApplier { let mut pending_data_count: u32 = 0; for data_id in &action_receipt.input_data_ids { storage.modify(|state_update| { - if get_received_data(state_update, account_id, *data_id) + if !has_received_data(state_update, account_id, *data_id) .expect("Genesis storage error") - .is_none() { pending_data_count += 1; set( diff --git a/core/store/src/lib.rs b/core/store/src/lib.rs index 9f1090f8b50..53b7184ce8c 100644 --- a/core/store/src/lib.rs +++ b/core/store/src/lib.rs @@ -703,6 +703,14 @@ pub fn get_received_data( get(trie, &TrieKey::ReceivedData { receiver_id: receiver_id.clone(), data_id }) } +pub fn has_received_data( + trie: &dyn TrieAccess, + receiver_id: &AccountId, + data_id: CryptoHash, +) -> Result { + trie.contains_key(&TrieKey::ReceivedData { receiver_id: receiver_id.clone(), data_id }) +} + pub fn set_postponed_receipt(state_update: &mut TrieUpdate, receipt: &Receipt) { let key = TrieKey::PostponedReceipt { receiver_id: receipt.receiver_id.clone(), diff --git a/runtime/runtime/src/lib.rs b/runtime/runtime/src/lib.rs index 481617b967b..822b76630d9 100644 --- a/runtime/runtime/src/lib.rs +++ b/runtime/runtime/src/lib.rs @@ -37,9 +37,9 @@ use near_primitives::utils::{ }; use near_primitives::version::{ProtocolFeature, ProtocolVersion}; use near_store::{ - get, get_account, get_postponed_receipt, get_received_data, remove_postponed_receipt, set, - set_account, set_delayed_receipt, set_postponed_receipt, set_received_data, PartialStorage, - StorageError, Trie, TrieChanges, TrieUpdate, + get, get_account, get_postponed_receipt, get_received_data, has_received_data, + remove_postponed_receipt, set, set_account, set_delayed_receipt, set_postponed_receipt, + set_received_data, PartialStorage, StorageError, Trie, TrieChanges, TrieUpdate, }; use near_store::{set_access_key, set_code}; use near_vm_runner::logic::types::PromiseResult; @@ -967,7 +967,7 @@ impl Runtime { // If not, then we will postpone this receipt for later. let mut pending_data_count: u32 = 0; for data_id in &action_receipt.input_data_ids { - if get_received_data(state_update, account_id, *data_id)?.is_none() { + if !has_received_data(state_update, account_id, *data_id)? { pending_data_count += 1; // The data for a given data_id is not available, so we save a link to this // receipt_id for the pending data_id into the state. From 441ab328c95b95551db01d17e397241d443da3f0 Mon Sep 17 00:00:00 2001 From: Simonas Kazlauskas Date: Fri, 26 Jan 2024 14:18:06 +0200 Subject: [PATCH 4/6] add a test --- core/store/src/trie/mod.rs | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/core/store/src/trie/mod.rs b/core/store/src/trie/mod.rs index 7f9ad859d7d..381207c98bc 100644 --- a/core/store/src/trie/mod.rs +++ b/core/store/src/trie/mod.rs @@ -1743,6 +1743,41 @@ mod tests { } } + #[test] + fn test_contains_key() { + let sid = ShardUId::single_shard(); + let bid = CryptoHash::default(); + let tries = TestTriesBuilder::new().with_flat_storage().build(); + let initial = vec![ + (vec![99, 44, 100, 58, 58, 49], Some(vec![1])), + (vec![99, 44, 100, 58, 58, 50], Some(vec![1])), + (vec![99, 44, 100, 58, 58, 50, 51], Some(vec![1])), + ]; + let root = test_populate_trie(&tries, &Trie::EMPTY_ROOT, sid, initial); + let trie = tries.get_trie_with_block_hash_for_shard(sid, root, &bid, false); + assert!(trie.has_flat_storage_chunk_view()); + assert!(trie.contains_key_mode(&[99, 44, 100, 58, 58, 49], KeyLookupMode::Trie).unwrap()); + assert!(trie + .contains_key_mode(&[99, 44, 100, 58, 58, 49], KeyLookupMode::FlatStorage) + .unwrap()); + assert!(!trie.contains_key_mode(&[99, 44, 100, 58, 58, 48], KeyLookupMode::Trie).unwrap()); + assert!(!trie + .contains_key_mode(&[99, 44, 100, 58, 58, 48], KeyLookupMode::FlatStorage) + .unwrap()); + let changes = vec![(vec![99, 44, 100, 58, 58, 49], None)]; + let root = test_populate_trie(&tries, &root, sid, changes); + let trie = tries.get_trie_with_block_hash_for_shard(sid, root, &bid, false); + assert!(trie.has_flat_storage_chunk_view()); + assert!(trie.contains_key_mode(&[99, 44, 100, 58, 58, 50], KeyLookupMode::Trie).unwrap()); + assert!(trie + .contains_key_mode(&[99, 44, 100, 58, 58, 50], KeyLookupMode::FlatStorage) + .unwrap()); + assert!(!trie + .contains_key_mode(&[99, 44, 100, 58, 58, 49], KeyLookupMode::FlatStorage) + .unwrap()); + assert!(!trie.contains_key_mode(&[99, 44, 100, 58, 58, 49], KeyLookupMode::Trie).unwrap()); + } + #[test] fn test_equal_leafs() { let initial = vec![ From 0b813a563bef5d3637962c3807afa81994858156 Mon Sep 17 00:00:00 2001 From: Simonas Kazlauskas Date: Fri, 26 Jan 2024 14:19:20 +0200 Subject: [PATCH 5/6] handle flat storage --- core/store/src/trie/mod.rs | 38 ++++++++++++++++++++++++++++++++------ 1 file changed, 32 insertions(+), 6 deletions(-) diff --git a/core/store/src/trie/mod.rs b/core/store/src/trie/mod.rs index 381207c98bc..e93170293c1 100644 --- a/core/store/src/trie/mod.rs +++ b/core/store/src/trie/mod.rs @@ -1370,14 +1370,40 @@ impl Trie { /// This method is guaranteed to not inspect the value stored for this key, which would /// otherwise have potential gas cost implications. pub fn contains_key(&self, key: &[u8]) -> Result { - let charge_gas_for_trie_node_access = self.charge_gas_for_trie_node_access; + self.contains_key_mode(key, KeyLookupMode::FlatStorage) + } + + /// Check if the column contains a value with the given `key`. + /// + /// This method is guaranteed to not inspect the value stored for this key, which would + /// otherwise have potential gas cost implications. + pub fn contains_key_mode(&self, key: &[u8], mode: KeyLookupMode) -> Result { + let charge_gas_for_trie_node_access = + mode == KeyLookupMode::Trie || self.charge_gas_for_trie_node_access; if self.memtries.is_some() { - Ok(self.lookup_from_memory(key, charge_gas_for_trie_node_access, |_| ())?.is_some()) - } else { - Ok(self - .lookup_from_state_column(NibbleSlice::new(key), charge_gas_for_trie_node_access)? - .is_some()) + return Ok(self + .lookup_from_memory(key, charge_gas_for_trie_node_access, |_| ())? + .is_some()); } + + 'flat: { + let KeyLookupMode::FlatStorage = mode else { break 'flat }; + let Some(flat_storage_chunk_view) = &self.flat_storage_chunk_view else { break 'flat }; + let value = flat_storage_chunk_view.contains_value(key)?; + if self.recorder.is_some() { + // If recording, we need to look up in the trie as well to record the trie nodes, + // as they are needed to prove the value. Also, it's important that this lookup + // is done even if the key was not found, because intermediate trie nodes may be + // needed to prove the non-existence of the key. + let value_ref_from_trie = + self.lookup_from_state_column(NibbleSlice::new(key), false)?; + debug_assert_eq!(&value_ref_from_trie.is_some(), &value); + } + } + + Ok(self + .lookup_from_state_column(NibbleSlice::new(key), charge_gas_for_trie_node_access)? + .is_some()) } /// Retrieves an `OptimizedValueRef`` for the given key. See `OptimizedValueRef`. From fb68b79aafd86b804d0395c1a4ad22493f9cc817 Mon Sep 17 00:00:00 2001 From: Simonas Kazlauskas Date: Mon, 29 Jan 2024 17:07:12 +0200 Subject: [PATCH 6/6] fixup rebase --- core/store/src/trie/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/store/src/trie/mod.rs b/core/store/src/trie/mod.rs index e93170293c1..912df1f4866 100644 --- a/core/store/src/trie/mod.rs +++ b/core/store/src/trie/mod.rs @@ -1389,7 +1389,7 @@ impl Trie { 'flat: { let KeyLookupMode::FlatStorage = mode else { break 'flat }; let Some(flat_storage_chunk_view) = &self.flat_storage_chunk_view else { break 'flat }; - let value = flat_storage_chunk_view.contains_value(key)?; + let value = flat_storage_chunk_view.contains_key(key)?; if self.recorder.is_some() { // If recording, we need to look up in the trie as well to record the trie nodes, // as they are needed to prove the value. Also, it's important that this lookup