Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

refactor: move nodes counter to the caching storage #6315

Merged
merged 34 commits into from
Mar 9, 2022
Merged

refactor: move nodes counter to the caching storage #6315

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
34 commits
Select commit Hold shift + click to select a range
a71e23b
move counter and implement test
Feb 17, 2022
3b86172
compilation fix
Feb 17, 2022
52d5c59
compilation fix
Feb 17, 2022
7c587df
compilation fix
Feb 17, 2022
df07bb9
compilation fix
Feb 17, 2022
c9db0d1
compilation fix
Feb 17, 2022
ffab524
Update core/store/src/trie/trie_storage.rs
Longarithm Feb 17, 2022
6b5327d
apply suggestions
Feb 17, 2022
c06f983
Merge branch 'master' into move-counter
Longarithm Feb 17, 2022
72c0be1
add test with repeated values
Feb 17, 2022
ad2fc60
merge with master
Feb 18, 2022
0517369
cargo fmt fix
Feb 18, 2022
5777028
add test contract rw cost test
Feb 22, 2022
d168d01
compilation fix
Feb 22, 2022
6f3c509
compilation fix
Feb 22, 2022
31e0471
compilation fix
Feb 22, 2022
3eabc3b
compilation fix
Feb 22, 2022
52d7bda
compilation fix
Feb 22, 2022
6f4c37d
compilation fix
Feb 22, 2022
2c1ae5b
cost name fix
Feb 22, 2022
b60fcf0
set value for test
Feb 22, 2022
a66b203
Merge branch 'master' into move-counter
Longarithm Feb 23, 2022
63a700b
debug print
Feb 23, 2022
0336b32
debug
Feb 23, 2022
314c3f3
debug
Feb 23, 2022
7aaf621
debug
Feb 23, 2022
9017bd4
debug
Feb 23, 2022
de97807
debug
Feb 23, 2022
118835d
debug
Feb 23, 2022
c4744a7
remove debug
Feb 23, 2022
df58c0f
remove debug
Feb 23, 2022
bb6ecba
remove debug
Feb 24, 2022
ab55083
remove default get_touched_nodes_count impl
Mar 9, 2022
09fd771
Merge branch 'master' into move-counter
Longarithm Mar 9, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion core/store/src/test_utils.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,7 @@ pub fn gen_larger_changes(rng: &mut impl Rng, max_size: usize) -> Vec<(Vec<u8>,
}

pub(crate) fn simplify_changes(
changes: &Vec<(Vec<u8>, Option<Vec<u8>>)>,
changes: &[(Vec<u8>, Option<Vec<u8>>)],
) -> Vec<(Vec<u8>, Option<Vec<u8>>)> {
let mut state: HashMap<Vec<u8>, Vec<u8>> = HashMap::new();
for (key, value) in changes.iter() {
Expand Down
3 changes: 2 additions & 1 deletion core/store/src/trie/iterator.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -285,7 +285,7 @@ impl<'a> TrieIterator<'a> {
}
IterStep::Continue => {}
IterStep::Value(hash) => {
self.trie.retrieve_raw_bytes(&hash)?;
self.trie.storage.retrieve_raw_bytes(&hash)?;
nodes_list.push(TrieTraversalItem {
hash,
key: self.has_value().then(|| self.key()),
Expand Down Expand Up @@ -322,6 +322,7 @@ impl<'a> Iterator for TrieIterator<'a> {
IterStep::Value(hash) => {
return Some(
self.trie
.storage
.retrieve_raw_bytes(&hash)
.map(|value| (self.key(), value.to_vec())),
)
Expand Down
29 changes: 11 additions & 18 deletions core/store/src/trie/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@ use std::cmp::Ordering;
use std::collections::HashMap;
use std::fmt;
use std::io::{Cursor, Read, Write};
use std::sync::Arc;

use borsh::{BorshDeserialize, BorshSerialize};
use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
Expand All @@ -18,10 +17,8 @@ use crate::trie::insert_delete::NodesStorage;
use crate::trie::iterator::TrieIterator;
use crate::trie::nibble_slice::NibbleSlice;
pub use crate::trie::shard_tries::{KeyForStateChanges, ShardTries, WrappedTrieChanges};
use crate::trie::trie_storage::{
TouchedNodesCounter, TrieMemoryPartialStorage, TrieRecordingStorage, TrieStorage,
};
pub(crate) use crate::trie::trie_storage::{TrieCache, TrieCachingStorage};
use crate::trie::trie_storage::{TrieMemoryPartialStorage, TrieRecordingStorage, TrieStorage};
use crate::StorageError;

mod insert_delete;
Expand Down Expand Up @@ -408,7 +405,6 @@ impl RawTrieNodeWithSize {

pub struct Trie {
pub(crate) storage: Box<dyn TrieStorage>,
pub counter: TouchedNodesCounter,
}

/// Stores reference count change for some key-value pair in DB.
Expand Down Expand Up @@ -472,7 +468,7 @@ pub struct ApplyStatePartResult {

impl Trie {
pub fn new(store: Box<dyn TrieStorage>, _shard_uid: ShardUId) -> Self {
Trie { storage: store, counter: TouchedNodesCounter::default() }
Trie { storage: store }
}

pub fn recording_reads(&self) -> Self {
Expand All @@ -483,7 +479,7 @@ impl Trie {
shard_uid: storage.shard_uid,
recorded: RefCell::new(Default::default()),
};
Trie { storage: Box::new(storage), counter: TouchedNodesCounter::default() }
Trie { storage: Box::new(storage) }
}

pub fn empty_root() -> StateRoot {
Expand All @@ -506,7 +502,6 @@ impl Trie {
recorded_storage,
visited_nodes: Default::default(),
}),
counter: TouchedNodesCounter::default(),
}
}

Expand Down Expand Up @@ -578,7 +573,6 @@ impl Trie {
if *hash == Trie::empty_root() {
Ok(memory.store(TrieNodeWithSize::empty()))
} else {
self.counter.increment();
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we have any protocol changes here? In particular, I see there's self.storage.retrieve_raw_bytes(hash)?; in delete_value, and that now would increment the counter.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great question!

From what I see, all logic leading to incrementing counter here happens in Trie::update which itself is called after all tx/receipt execution. I feel that some counter increments makes us think that we charge for them somewhere, although we don't do it in reality. So it's better to remove them. But I also will be happy if someone double-check this.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a way to test this?

Copy link
Member Author

@Longarithm Longarithm Feb 22, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think that what exactly we need to test is that gas costs were not affected by this change - because they are part of protocol.

The similar purpose is served by the test group next to Runtime implementation. In particular, they check gas costs after Runtime::apply() call, which include Trie::update call.

However, we currently have only one test applying funcall txs - test_apply_deficit_gas_for_function_call_partial. So I will add a new test there which will test a funcall which touches contract storage. This should help in the future - if we accidentally change counter, the test must fail.

UPD: added a test test_contract_read_write_cost

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Longarithm I still don't understand why the change in delete_value won't cause any change in gas counting. Could you elaborate more on that?

One way to test this manually, is maybe run state viewer apply chain on some mainnet blocks to verify that gas counting doesn't change.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@mzhangmzz delete_value remains unchanged here, it is a change in move_node_to_mutable.

Anyway it won't cause changes in gas counting, because these methods are called only on Trie::update, which happens after tx/receipt processing as I said in the second comment. Having the increment I am removing actually didn't make much sense.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think delete_value will change here, because previously it only calls self.storage.retrieve_raw_bytes(hash)? without incrementing the counter. After your change, since you moved the counter incrementing logic inside storage.retrieve_raw_bytes, delete_value will increase the counter.

I'm not super familiar with the gas counting logic, so I'm not sure why it is called in Trie::update won't affect gas counting.

Copy link
Member Author

@Longarithm Longarithm Feb 24, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, you are right.

I'm not sure why it is called in Trie::update won't affect gas counting.

I've looked over all Trie::update calls. It is called for four purposes:

  • tests
  • splitting states (add_values_to_split_states_impl)
  • applying genesis (finalize_genesis)
  • in TrieUpdate::finalize

The latter function is called for similar purposes - I think it doesn't make much sense to list all them. Also it is called in the end of Runtime::apply to apply changes to trie, when all tx/receipts are already executed and burnt/used gas is computed.

Moreover, trie counter is used for gas counting only during function call execution, which is triggered only by process_receipt, this is another reason why Trie::update change doesn't affect it.

Sorry I didn't explain it before properly. On my side, it was a fast look over CLion usages, so I thought it didn't worth attention.

Also, I could leave the incrementation logic as is. But it would make the implementation more ugly, because by default TrieStorage shouldn't know about counter.

let bytes = self.storage.retrieve_raw_bytes(hash)?;
match RawTrieNodeWithSize::decode(&bytes) {
Ok(value) => {
Expand All @@ -602,7 +596,7 @@ impl Trie {
if *hash == Trie::empty_root() {
return Ok(TrieNodeWithSize::empty());
}
let bytes = self.retrieve_raw_bytes(hash)?;
let bytes = self.storage.retrieve_raw_bytes(hash)?;
match RawTrieNodeWithSize::decode(&bytes) {
Ok(value) => Ok(TrieNodeWithSize::from_raw(value)),
Err(_) => Err(StorageError::StorageInconsistentState(format!(
Expand All @@ -612,16 +606,11 @@ impl Trie {
}
}

pub(crate) fn retrieve_raw_bytes(&self, hash: &CryptoHash) -> Result<Arc<[u8]>, StorageError> {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@matklad - more Rust question:

in such case - should we 'keep' this method (and have it simply call self.storage.retrieve_raw_bytes) or should we allow other classes to call XX.storage.retrieve_raw_bytes (as Alex did in this PR)?

For me - one advantage of keeping this method - is that we keep better 'layers' - and avoid having other classes 'peek' into Trie) - but I wonder if it has some performance implications

Copy link
Contributor

@matklad matklad Mar 15, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@mm-near there won't be performance implications here. In general, it's a pretty safe assumption in Rust that linguistic abstractions don't introduce overhead -- static nature of the language and unboxed memory layout guarantee that compiler "sees through" any abstractions.

Regarding layering, my general observation is that layering is good only when it is enforced: it's better to focus on specific places where layering is important, rather than to try to enforce it recursively throughout the stack. There should be places where there's explicitly no layers.

For Rust, the best enforcement is aligning layers with crate boundaries: that's why I am relatively picky about pub vs pub(crate), but relatively indifferent between pub(crate) and /*private*/.

As an aside, I think we can make storage private rather pub(crate) if we put trie tests into submodule of the trie module. We only need pub(crate) for tests, but "pub(crate) for tests" is a bit of oxymoron -- pub(crate) guarantees that this is for unit tests, and unit tests can always be made a submodule (in Rust, submodule see the full contents of partent modules)

self.counter.increment();
self.storage.retrieve_raw_bytes(hash)
}

pub fn retrieve_root_node(&self, root: &StateRoot) -> Result<StateRootNode, StorageError> {
if *root == Trie::empty_root() {
return Ok(StateRootNode::empty());
}
let data = self.retrieve_raw_bytes(root)?;
let data = self.storage.retrieve_raw_bytes(root)?;
match RawTrieNodeWithSize::decode(&data) {
Ok(value) => {
let memory_usage = TrieNodeWithSize::from_raw(value).memory_usage;
Expand All @@ -645,7 +634,7 @@ impl Trie {
if hash == Trie::empty_root() {
return Ok(None);
}
let bytes = self.retrieve_raw_bytes(&hash)?;
let bytes = self.storage.retrieve_raw_bytes(&hash)?;
let node = RawTrieNodeWithSize::decode(&bytes).map_err(|_| {
StorageError::StorageInconsistentState("RawTrieNode decode failed".to_string())
})?;
Expand Down Expand Up @@ -701,7 +690,7 @@ impl Trie {
pub fn get(&self, root: &CryptoHash, key: &[u8]) -> Result<Option<Vec<u8>>, StorageError> {
match self.get_ref(root, key)? {
Some((_length, hash)) => {
self.retrieve_raw_bytes(&hash).map(|bytes| Some(bytes.to_vec()))
self.storage.retrieve_raw_bytes(&hash).map(|bytes| Some(bytes.to_vec()))
}
None => Ok(None),
}
Expand Down Expand Up @@ -761,6 +750,10 @@ impl Trie {
pub fn iter<'a>(&'a self, root: &CryptoHash) -> Result<TrieIterator<'a>, StorageError> {
TrieIterator::new(self, root)
}

pub fn get_touched_nodes_count(&self) -> u64 {
self.storage.get_touched_nodes_count()
}
}

#[cfg(test)]
Expand Down
2 changes: 1 addition & 1 deletion core/store/src/trie/state_parts.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -218,7 +218,7 @@ impl Trie {
let mut map = HashMap::new();
let mut contract_codes = Vec::new();
for TrieTraversalItem { hash, key } in trie_traversal_items {
let value = trie.retrieve_raw_bytes(&hash)?;
let value = trie.storage.retrieve_raw_bytes(&hash)?;
map.entry(hash).or_insert_with(|| (value.to_vec(), 0)).1 += 1;
if let Some(trie_key) = key {
if is_contract_code_key(&trie_key) {
Expand Down
46 changes: 28 additions & 18 deletions core/store/src/trie/trie_storage.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
use std::collections::{HashMap, HashSet};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

use near_primitives::hash::CryptoHash;
Expand All @@ -9,7 +8,7 @@ use crate::trie::POISONED_LOCK_ERR;
use crate::{ColState, StorageError, Store};
use lru::LruCache;
use near_primitives::shard_layout::ShardUId;
use std::cell::RefCell;
use std::cell::{Cell, RefCell};
use std::io::ErrorKind;

#[derive(Clone)]
Expand Down Expand Up @@ -40,6 +39,12 @@ impl TrieCache {
}
}
}

#[cfg(test)]
pub(crate) fn len(&self) -> usize {
let guard = self.0.lock().expect(POISONED_LOCK_ERR);
guard.len()
}
}

pub trait TrieStorage {
Expand All @@ -59,10 +64,13 @@ pub trait TrieStorage {
fn as_partial_storage(&self) -> Option<&TrieMemoryPartialStorage> {
None
}

fn get_touched_nodes_count(&self) -> u64;
}

/// Records every value read by retrieve_raw_bytes.
/// Used for obtaining state parts (and challenges in the future).
/// TODO (#6316): implement proper nodes counting logic as in TrieCachingStorage
pub struct TrieRecordingStorage {
pub(crate) store: Store,
pub(crate) shard_uid: ShardUId,
Expand Down Expand Up @@ -90,6 +98,10 @@ impl TrieStorage for TrieRecordingStorage {
fn as_recording_storage(&self) -> Option<&TrieRecordingStorage> {
Some(self)
}

fn get_touched_nodes_count(&self) -> u64 {
unimplemented!();
}
}

/// Storage for validating recorded partial storage.
Expand All @@ -114,6 +126,10 @@ impl TrieStorage for TrieMemoryPartialStorage {
fn as_partial_storage(&self) -> Option<&TrieMemoryPartialStorage> {
Some(self)
}

fn get_touched_nodes_count(&self) -> u64 {
unimplemented!();
}
}

/// Maximum number of cache entries.
Expand All @@ -135,11 +151,13 @@ pub struct TrieCachingStorage {
pub(crate) store: Store,
pub(crate) cache: TrieCache,
pub(crate) shard_uid: ShardUId,

pub(crate) counter: Cell<u64>,
}

impl TrieCachingStorage {
pub fn new(store: Store, cache: TrieCache, shard_uid: ShardUId) -> TrieCachingStorage {
TrieCachingStorage { store, cache, shard_uid }
TrieCachingStorage { store, cache, shard_uid, counter: Cell::new(0u64) }
}

pub(crate) fn get_shard_uid_and_hash_from_key(
Expand All @@ -162,10 +180,15 @@ impl TrieCachingStorage {
key[8..].copy_from_slice(hash.as_ref());
key
}

fn inc_counter(&self) {
self.counter.set(self.counter.get() + 1);
}
}

impl TrieStorage for TrieCachingStorage {
fn retrieve_raw_bytes(&self, hash: &CryptoHash) -> Result<Arc<[u8]>, StorageError> {
self.inc_counter();
let mut guard = self.cache.0.lock().expect(POISONED_LOCK_ERR);
if let Some(val) = guard.get(hash) {
Ok(val.clone())
Expand All @@ -191,21 +214,8 @@ impl TrieStorage for TrieCachingStorage {
fn as_caching_storage(&self) -> Option<&TrieCachingStorage> {
Some(self)
}
}

/// Runtime counts the number of touched trie nodes for the purpose of gas calculation.
/// Trie increments it on every call to TrieStorage::retrieve_raw_bytes()
#[derive(Default)]
pub struct TouchedNodesCounter {
counter: AtomicU64,
}

impl TouchedNodesCounter {
pub fn increment(&self) {
self.counter.fetch_add(1, Ordering::SeqCst);
}

pub fn get(&self) -> u64 {
self.counter.load(Ordering::SeqCst)
fn get_touched_nodes_count(&self) -> u64 {
self.counter.get()
}
}
80 changes: 79 additions & 1 deletion core/store/src/trie/trie_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,10 @@ impl TrieStorage for IncompletePartialStorage {
// Make sure it's not called - it pretends to be PartialStorage but is not
unimplemented!()
}

fn get_touched_nodes_count(&self) -> u64 {
unimplemented!();
}
}

fn setup_storage<F, Out>(trie: Rc<Trie>, test: &mut F) -> (PartialStorage, Out)
Expand All @@ -75,7 +79,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), counter: Default::default() };
let trie = Trie { storage: Box::new(storage) };
let expected_result =
if i < size { Err(&StorageError::TrieNodeMissing) } else { Ok(&expected) };
assert_eq!(test(Rc::new(trie)).as_ref(), expected_result);
Expand Down Expand Up @@ -127,3 +131,77 @@ fn test_reads_with_incomplete_storage() {
}
}
}

#[cfg(test)]
mod caching_storage_tests {
use super::*;
use crate::test_utils::create_tries;
use crate::trie::nibble_slice::NibbleSlice;

fn create_trie_key(nibbles: &[u8]) -> Vec<u8> {
NibbleSlice::encode_nibbles(&nibbles, false).into_vec()
}

fn create_trie(items: &[(Vec<u8>, Option<Vec<u8>>)]) -> (Rc<Trie>, CryptoHash) {
let tries = create_tries();
let shard_uid = ShardUId { version: 1, shard_id: 0 };
let trie = tries.get_trie_for_shard(shard_uid);
let trie = Rc::new(trie);
let state_root = Trie::empty_root();
let trie_changes = simplify_changes(&items);
let state_root = test_populate_trie(&tries, &state_root, shard_uid, trie_changes.clone());
(trie, state_root)
}

// Get values corresponding to keys one by one, returning vector of numbers of touched nodes for each `get`.
fn get_touched_nodes_numbers(
trie: Rc<Trie>,
state_root: CryptoHash,
items: &[(Vec<u8>, Option<Vec<u8>>)],
) -> Vec<u64> {
items
.iter()
.map(|(key, value)| {
let initial_counter = trie.get_touched_nodes_count();
let got_value = trie.get(&state_root, key).unwrap();
assert_eq!(*value, got_value);
trie.get_touched_nodes_count() - initial_counter
})
.collect()
}

// Test nodes counter and trie cache size on the sample of trie items.
#[test]
fn count_touched_nodes() {
// For keys with nibbles [000, 011, 100], we expect 6 touched nodes to get value for the first key 000:
// Extension -> Branch -> Branch -> Leaf plus retrieving the value by its hash. In total
// there will be 9 distinct nodes, because 011 and 100 both add one Leaf and value.
let trie_items = vec![
(create_trie_key(&vec![0, 0, 0]), Some(vec![0])),
(create_trie_key(&vec![0, 1, 1]), Some(vec![1])),
(create_trie_key(&vec![1, 0, 0]), Some(vec![2])),
];
let (trie, state_root) = create_trie(&trie_items);
assert_eq!(get_touched_nodes_numbers(trie.clone(), state_root, &trie_items), vec![5, 5, 4]);

let storage = trie.storage.as_caching_storage().unwrap();
assert_eq!(storage.cache.len(), 9);
}

// Check that same values are stored in the same trie node.
#[test]
fn count_touched_nodes_repeated_values() {
// For these keys there will be 5 nodes with distinct hashes, because each path looks like
// Extension([0, 0]) -> Branch -> Leaf([48/49]) -> value.
// TODO: explain the exact values in path items here
let trie_items = vec![
(create_trie_key(&vec![0, 0]), Some(vec![1])),
(create_trie_key(&vec![1, 1]), Some(vec![1])),
];
let (trie, state_root) = create_trie(&trie_items);
assert_eq!(get_touched_nodes_numbers(trie.clone(), state_root, &trie_items), vec![4, 4]);

let storage = trie.storage.as_caching_storage().unwrap();
assert_eq!(storage.cache.len(), 5);
}
}
2 changes: 1 addition & 1 deletion core/store/src/trie/update.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ impl<'a> TrieUpdateValuePtr<'a> {
match self {
TrieUpdateValuePtr::MemoryRef(value) => Ok((*value).clone()),
TrieUpdateValuePtr::HashAndSize(trie, _, hash) => {
trie.retrieve_raw_bytes(hash).map(|bytes| bytes.to_vec())
trie.storage.retrieve_raw_bytes(hash).map(|bytes| bytes.to_vec())
}
}
}
Expand Down
Loading