From a679d95de409b93f800ab92f10aee01af3f36be4 Mon Sep 17 00:00:00 2001 From: Simonas Kazlauskas Date: Fri, 27 Sep 2024 14:41:32 +0300 Subject: [PATCH] Add additional mechanisms to introspect state witnesses In particular this code is what I had used to inspect the contents of the state witnesses and the contribution of trie's contents to the state witness size. This may or may not be useful, so am submitting this as a PR. --- core/primitives/src/challenge.rs | 3 +- .../src/stateless_validation/state_witness.rs | 2 +- tools/state-viewer/src/latest_witnesses.rs | 125 +++++++++++++++++- 3 files changed, 122 insertions(+), 8 deletions(-) diff --git a/core/primitives/src/challenge.rs b/core/primitives/src/challenge.rs index 5f9a3dbeb8c..bae0b79f0ec 100644 --- a/core/primitives/src/challenge.rs +++ b/core/primitives/src/challenge.rs @@ -31,7 +31,8 @@ impl Debug for PartialState { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { PartialState::TrieValues(values) => { - f.write_str(&format!("{} trie values", values.len())) + let size = bytesize::ByteSize(values.iter().map(|v| v.len() as u64).sum()); + f.write_str(&format!("{} trie values sized {}", values.len(), size)) } } } diff --git a/core/primitives/src/stateless_validation/state_witness.rs b/core/primitives/src/stateless_validation/state_witness.rs index cf82101f036..600fb5c029e 100644 --- a/core/primitives/src/stateless_validation/state_witness.rs +++ b/core/primitives/src/stateless_validation/state_witness.rs @@ -145,7 +145,7 @@ pub struct ChunkStateWitness { pub new_transactions: Vec, pub new_transactions_validation_state: PartialState, // TODO(stateless_validation): Deprecate once we send state witness in parts. - signature_differentiator: SignatureDifferentiator, + pub signature_differentiator: SignatureDifferentiator, } impl ChunkStateWitness { diff --git a/tools/state-viewer/src/latest_witnesses.rs b/tools/state-viewer/src/latest_witnesses.rs index ec3e4edd54d..21ca5b65796 100644 --- a/tools/state-viewer/src/latest_witnesses.rs +++ b/tools/state-viewer/src/latest_witnesses.rs @@ -1,14 +1,18 @@ +use std::collections::BTreeMap; use std::path::{Path, PathBuf}; use std::rc::Rc; +use borsh::BorshDeserialize as _; use near_chain::runtime::NightshadeRuntime; use near_chain::stateless_validation::processing_tracker::ProcessingDoneTracker; use near_chain::{Chain, ChainGenesis, ChainStore, DoomslugThresholdMode}; use near_epoch_manager::shard_tracker::{ShardTracker, TrackedConfig}; use near_epoch_manager::EpochManager; +use near_primitives::challenge::PartialState; +use near_primitives::hash::{hash, CryptoHash}; use near_primitives::stateless_validation::state_witness::ChunkStateWitness; use near_primitives::types::EpochId; -use near_store::Store; +use near_store::{RawTrieNodeWithSize, Store}; use near_time::Clock; use nearcore::NearConfig; use nearcore::NightshadeRuntimeExt; @@ -48,13 +52,76 @@ pub struct DumpWitnessesCmd { #[derive(clap::Subcommand)] enum DumpWitnessesMode { - /// Pretty-print on screen using the "{:#?}" formatting. - Pretty, + /// Pretty-print the contents on screen. + Pretty { + #[clap(long)] + decode_tries: bool, + }, /// Saves the raw &[u8] of each witness to the given directory. Binary { output_dir: PathBuf }, } impl DumpWitnessesCmd { + fn pretty_print_partial_state(&self, partial_state: &PartialState) { + let PartialState::TrieValues(nodes) = partial_state; + let mut nmap = nodes.iter().map(|n| (hash(&n), (None, &**n))).collect::>(); + fn compute_size(nmap: &mut BTreeMap, &[u8])>, h: &CryptoHash) { + let Some((sz, val)) = nmap.get(h) else { return }; + if sz.is_some() { + return; + } + let self_size = val.len() as u64; + let Ok(node) = RawTrieNodeWithSize::try_from_slice(val) else { + nmap.get_mut(h).unwrap().0.replace(self_size); + return; + }; + let size = match node.node { + near_store::RawTrieNode::Leaf(_, r) => r.length as u64, + near_store::RawTrieNode::BranchNoValue(ch) => ch + .iter() + .map(|(_, c)| { + compute_size(nmap, c); + nmap.get(c).map(|v| v.0.unwrap()).unwrap_or(0) + }) + .sum(), + near_store::RawTrieNode::BranchWithValue(v, ch) => { + ch.iter() + .map(|(_, c)| { + compute_size(nmap, c); + nmap.get(c).map(|v| v.0.unwrap()).unwrap_or(0) + }) + .sum::() + + v.length as u64 + } + near_store::RawTrieNode::Extension(_, c) => { + compute_size(nmap, &c); + nmap.get(&c).map(|v| v.0.unwrap()).unwrap_or(0) + } + }; + nmap.get_mut(h).unwrap().0.replace(self_size + size); + } + + for n in nmap.keys().copied().collect::>() { + compute_size(&mut nmap, &n); + } + + for node in nodes { + let hash = hash(node); + let Ok(node) = RawTrieNodeWithSize::try_from_slice(node) else { + println!(" {}: Not RawTrieNodeWithSize (size: {})", hash, node.len()); + continue; + }; + + println!( + " {}: .memory_usage: {}, .actual_cumulative_memory: {}, .node: {:?}", + hash, + bytesize::ByteSize(node.memory_usage), + bytesize::ByteSize(nmap.get(&hash).unwrap().0.unwrap()), + node.node, + ); + } + } + pub(crate) fn run(&self, near_config: NearConfig, store: Store) { let chain_store = Rc::new(ChainStore::new(store, near_config.genesis.config.genesis_height, false)); @@ -77,9 +144,55 @@ impl DumpWitnessesCmd { witness.epoch_id ); match self.mode { - DumpWitnessesMode::Pretty => { - println!("{:#?}", witness); - println!(""); + DumpWitnessesMode::Pretty { decode_tries } => { + let ChunkStateWitness { + chunk_producer, + epoch_id, + chunk_header, + main_state_transition, + source_receipt_proofs, + applied_receipts_hash, + transactions, + implicit_transitions, + new_transactions, + new_transactions_validation_state, + signature_differentiator, + } = witness; + println!("chunk_producer: {:#?}", chunk_producer); + println!("epoch_id: {:#?}", epoch_id); + println!("chunk_header: {:#?}", chunk_header); + println!( + "main_state_transition.block_hash: {:#?}", + main_state_transition.block_hash + ); + println!( + "main_state_transition.post_state_root: {:#?}", + main_state_transition.post_state_root + ); + if decode_tries { + println!("main_state_transition.base_state: (continued)",); + self.pretty_print_partial_state(&main_state_transition.base_state); + } else { + println!( + "main_state_transition.base_state: {:#?}", + main_state_transition.base_state + ); + } + println!("source_receipt_proofs: {:#?}", source_receipt_proofs); + println!("applied_receipts_hash: {:#?}", applied_receipts_hash); + println!("transactions: {:#?}", transactions); + println!("implicit_transitions: {:#?}", implicit_transitions); + println!("new_transactions: {:#?}", new_transactions); + if decode_tries { + println!("new_transactions_validation_state: (continued)",); + self.pretty_print_partial_state(new_transactions_validation_state); + } else { + println!( + "new_transactions_validation_state: {:#?}", + new_transactions_validation_state + ); + } + println!("signature_differentiator: {:#?}", signature_differentiator); } DumpWitnessesMode::Binary { ref output_dir } => { let file_name = format!(