feat: support hot/cold store in state viewer and add view-trie2 comma…

…nd (#8311)

Added --store-temperature to view-state and implemented a new, pretty, output format for view-steate view-trie.

chain/network/src/peer_manager/peer_store/mod.rs

@@ -7,11 +7,13 @@ use anyhow::bail;
 use im::hashmap::Entry;
 use im::{HashMap, HashSet};
 use near_primitives::network::PeerId;
+use near_store::db::Database;
 use parking_lot::Mutex;
 use rand::seq::IteratorRandom;
 use rand::thread_rng;
 use std::net::SocketAddr;
 use std::ops::Not;
+use std::sync::Arc;
 
 #[cfg(test)]
 mod testonly;
@@ -592,12 +594,12 @@ impl PeerStore {
 }
 
 /// Public method used to iterate through all peers stored in the database.
-pub fn iter_peers_from_store<F>(store: near_store::NodeStorage, f: F)
+pub fn iter_peers_from_store<F>(db: Arc<dyn Database>, f: F)
 where
     F: Fn((PeerId, KnownPeerState)),
 {
-    let db = store.into_inner(near_store::Temperature::Hot);
-    for x in crate::store::Store::from(db).list_peer_states().unwrap() {
+    let store = crate::store::Store::from(db);
+    for x in store.list_peer_states().unwrap() {
         f(x)
     }
 }

core/store/src/lib.rs

@@ -2,6 +2,7 @@ use std::fs::File;
 use std::io::{BufReader, BufWriter, Read, Write};
 use std::marker::PhantomData;
 use std::path::Path;
+use std::str::FromStr;
 use std::sync::Arc;
 use std::{fmt, io};
 
@@ -69,6 +70,20 @@ pub enum Temperature {
     Cold,
 }
 
+impl FromStr for Temperature {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, String> {
+        let s = s.to_lowercase();
+        match s.as_str() {
+            "hot" => Ok(Temperature::Hot),
+            #[cfg(feature = "cold_store")]
+            "cold" => Ok(Temperature::Cold),
+            _ => Err(String::from(format!("invalid temperature string {s}"))),
+        }
+    }
+}
+
 /// Node’s storage holding chain and all other necessary data.
 ///
 /// The eventual goal is to implement cold storage at which point this structure

core/store/src/trie/iterator.rs

@@ -4,13 +4,17 @@ use crate::trie::nibble_slice::NibbleSlice;
 use crate::trie::{TrieNode, TrieNodeWithSize, ValueHandle};
 use crate::{StorageError, Trie};
 
+/// Crumb is a piece of trie iteration state. It describes a node on the trail and processing status of that node.
 #[derive(Debug)]
 struct Crumb {
     node: TrieNodeWithSize,
     status: CrumbStatus,
     prefix_boundary: bool,
 }
 
+/// The status of processing of a node during trie iteration.
+/// Each node is processed in the following order:
+/// Entering -> At -> AtChild(0) -> ... -> AtChild(15) -> Exiting
 #[derive(Clone, Copy, Eq, PartialEq, Debug)]
 pub(crate) enum CrumbStatus {
     Entering,
@@ -37,15 +41,27 @@ impl Crumb {
     }
 }
 
+/// Trie iteration is done using a stack based approach.
+/// There are two stacks that we track while iterating: the trail and the key_nibbles.
+/// The trail is a vector of trie nodes on the path from root node to the node that is
+/// currently being processed together with processing status - the Crumb.
+/// The key_nibbles is a vector of nibbles from the state root not to the node that is
+/// currently being processed.
+/// The trail and the key_nibbles may have different lengths e.g. an extension trie node
+/// will add only a single item to the trail but may add multiple nibbles to the key_nibbles.
 pub struct TrieIterator<'a> {
     trie: &'a Trie,
     trail: Vec<Crumb>,
     pub(crate) key_nibbles: Vec<u8>,
 
     /// If not `None`, a list of all nodes that the iterator has visited.
     visited_nodes: Option<Vec<std::sync::Arc<[u8]>>>,
+
+    /// Max depth of iteration.
+    max_depth: Option<usize>,
 }
 
+/// The TrieTiem is a tuple of (key, value) of the node.
 pub type TrieItem = (Vec<u8>, Vec<u8>);
 
 /// Item extracted from Trie during depth first traversal, corresponding to some Trie node.
@@ -59,12 +75,13 @@ pub struct TrieTraversalItem {
 impl<'a> TrieIterator<'a> {
     #![allow(clippy::new_ret_no_self)]
     /// Create a new iterator.
-    pub(super) fn new(trie: &'a Trie) -> Result<Self, StorageError> {
+    pub(super) fn new(trie: &'a Trie, max_depth: Option<usize>) -> Result<Self, StorageError> {
         let mut r = TrieIterator {
             trie,
             trail: Vec::with_capacity(8),
             key_nibbles: Vec::with_capacity(64),
             visited_nodes: None,
+            max_depth,
         };
         r.descend_into_node(&trie.root)?;
         Ok(r)
@@ -367,16 +384,24 @@ impl<'a> Iterator for TrieIterator<'a> {
     fn next(&mut self) -> Option<Self::Item> {
         loop {
             let iter_step = self.iter_step()?;
-            match iter_step {
-                IterStep::PopTrail => {
+
+            let can_process = match self.max_depth {
+                Some(max_depth) => self.key_nibbles.len() <= max_depth,
+                None => true,
+            };
+
+            match (iter_step, can_process) {
+                (IterStep::Continue, _) => {}
+                (IterStep::PopTrail, _) => {
                     self.trail.pop();
                 }
-                IterStep::Descend(hash) => match self.descend_into_node(&hash) {
+                // Skip processing the node if can process is false.
+                (_, false) => {}
+                (IterStep::Descend(hash), true) => match self.descend_into_node(&hash) {
                     Ok(_) => (),
                     Err(err) => return Some(Err(err)),
                 },
-                IterStep::Continue => {}
-                IterStep::Value(hash) => {
+                (IterStep::Value(hash), true) => {
                     return Some(
                         self.trie
                             .storage

core/store/src/trie/mod.rs

@@ -1,6 +1,8 @@
 use std::cell::RefCell;
 use std::collections::HashMap;
+use std::fmt::Write;
 use std::io::Read;
+use std::str;
 
 use borsh::{BorshDeserialize, BorshSerialize};
 use byteorder::{LittleEndian, ReadBytesExt};
@@ -12,6 +14,7 @@ pub use near_primitives::shard_layout::ShardUId;
 use near_primitives::state::ValueRef;
 #[cfg(feature = "protocol_feature_flat_state")]
 use near_primitives::state_record::is_delayed_receipt_key;
+use near_primitives::state_record::StateRecord;
 use near_primitives::trie_key::TrieKey;
 use near_primitives::types::{StateRoot, StateRootNode};
 
@@ -27,7 +30,6 @@ pub use crate::trie::trie_storage::{TrieCache, TrieCachingStorage, TrieDBStorage
 use crate::trie::trie_storage::{TrieMemoryPartialStorage, TrieRecordingStorage};
 use crate::{FlatStateDelta, StorageError};
 pub use near_primitives::types::TrieNodesCount;
-use std::fmt::Write;
 
 mod config;
 mod insert_delete;
@@ -556,6 +558,11 @@ pub struct ApplyStatePartResult {
     pub contract_codes: Vec<ContractCode>,
 }
 
+enum NodeOrValue {
+    Node(RawTrieNodeWithSize),
+    Value(std::sync::Arc<[u8]>),
+}
+
 impl Trie {
     pub const EMPTY_ROOT: StateRoot = StateRoot::new();
 
@@ -661,14 +668,15 @@ impl Trie {
     }
 
     // Prints the trie nodes starting from hash, up to max_depth depth.
+    // The node hash can be any node in the trie.
     pub fn print_recursive(&self, f: &mut dyn std::io::Write, hash: &CryptoHash, max_depth: u32) {
         match self.retrieve_raw_node_or_value(hash) {
-            Ok(Ok(_)) => {
+            Ok(NodeOrValue::Node(_)) => {
                 let mut prefix: Vec<u8> = Vec::new();
                 self.print_recursive_internal(f, hash, max_depth, &mut "".to_string(), &mut prefix)
                     .expect("write failed");
             }
-            Ok(Err(value_bytes)) => {
+            Ok(NodeOrValue::Value(value_bytes)) => {
                 writeln!(
                     f,
                     "Given node is a value. Len: {}, Data: {:?} ",
@@ -683,6 +691,38 @@ impl Trie {
         };
     }
 
+    // Prints the trie leaves starting from the state root node, up to max_depth depth.
+    // This method can only iterate starting from the root node and it only prints the
+    // leaf nodes but it shows output in more human friendly way.
+    pub fn print_recursive_leaves(&self, f: &mut dyn std::io::Write, max_depth: u32) {
+        let iter = match self.iter_with_max_depth(max_depth as usize) {
+            Ok(iter) => iter,
+            Err(err) => {
+                writeln!(f, "Error when getting the trie iterator: {}", err).expect("write failed");
+                return;
+            }
+        };
+        for node in iter {
+            let (key, value) = match node {
+                Ok((key, value)) => (key, value),
+                Err(err) => {
+                    writeln!(f, "Failed to iterate node with error: {err}").expect("write failed");
+                    continue;
+                }
+            };
+
+            // Try to parse the key in UTF8 which works only for the simplest keys (e.g. account),
+            // or get whitespace padding instead.
+            let key_string = match str::from_utf8(&key) {
+                Ok(value) => String::from(value),
+                Err(_) => " ".repeat(key.len()),
+            };
+            let state_record = StateRecord::from_raw_key_value(key.clone(), value);
+
+            writeln!(f, "{} {state_record:?}", key_string).expect("write failed");
+        }
+    }
+
     // Converts the list of Nibbles to a readable string.
     fn nibbles_to_string(&self, prefix: &[u8]) -> String {
         let (chunks, remainder) = stdx::as_chunks::<2, _>(prefix);
@@ -795,12 +835,14 @@ impl Trie {
     }
 
     // Similar to retrieve_raw_node but handles the case where there is a Value (and not a Node) in the database.
-    fn retrieve_raw_node_or_value(
-        &self,
-        hash: &CryptoHash,
-    ) -> Result<Result<RawTrieNodeWithSize, std::sync::Arc<[u8]>>, StorageError> {
+    // This method is not safe to be used in any real scenario as it can incorrectly interpret a value as a trie node.
+    // It's only provided as a convenience for debugging tools.
+    fn retrieve_raw_node_or_value(&self, hash: &CryptoHash) -> Result<NodeOrValue, StorageError> {
         let bytes = self.storage.retrieve_raw_bytes(hash)?;
-        Ok(RawTrieNodeWithSize::decode(&bytes).map_err(|_| bytes))
+        match RawTrieNodeWithSize::decode(&bytes) {
+            Ok(node) => Ok(NodeOrValue::Node(node)),
+            Err(_) => Ok(NodeOrValue::Value(bytes)),
+        }
     }
 
     fn move_node_to_mutable(
@@ -984,7 +1026,14 @@ impl Trie {
     }
 
     pub fn iter<'a>(&'a self) -> Result<TrieIterator<'a>, StorageError> {
-        TrieIterator::new(self)
+        TrieIterator::new(self, None)
+    }
+
+    pub fn iter_with_max_depth<'a>(
+        &'a self,
+        max_depth: usize,
+    ) -> Result<TrieIterator<'a>, StorageError> {
+        TrieIterator::new(self, Some(max_depth))
     }
 
     pub fn get_trie_nodes_count(&self) -> TrieNodesCount {

neard/src/cli.rs

@@ -94,7 +94,7 @@ impl NeardCmd {
 
             NeardSubCommand::StateViewer(cmd) => {
                 let mode = if cmd.readwrite { Mode::ReadWrite } else { Mode::ReadOnly };
-                cmd.subcmd.run(&home_dir, genesis_validation, mode);
+                cmd.subcmd.run(&home_dir, genesis_validation, mode, cmd.store_temperature);
             }
 
             NeardSubCommand::RecompressStorage(cmd) => {
@@ -131,6 +131,10 @@ pub(super) struct StateViewerCommand {
     /// In case an operation needs to write to caches, a read-write mode may be needed.
     #[clap(long, short = 'w')]
     readwrite: bool,
+    /// What store temperature should the state viewer open. Allowed values are hot and cold but
+    /// cold is only available when cold_store feature is enabled.
+    #[clap(long, short = 't', default_value = "hot")]
+    store_temperature: near_store::Temperature,
     #[clap(subcommand)]
     subcmd: StateViewerSubCommand,
 }

tools/cold-store/src/cli.rs

@@ -45,15 +45,15 @@ fn check_open(store: &NodeStorage) {
 
 fn print_heads(store: &NodeStorage) {
     println!(
-        "HOT HEAD is at {:?}",
+        "HOT HEAD is at {:#?}",
         store.get_store(Temperature::Hot).get_ser::<Tip>(DBCol::BlockMisc, HEAD_KEY)
     );
     println!(
-        "HOT FINAL HEAD is at {:?}",
+        "HOT FINAL HEAD is at {:#?}",
         store.get_store(Temperature::Hot).get_ser::<Tip>(DBCol::BlockMisc, FINAL_HEAD_KEY)
     );
     println!(
-        "COLD HEAD is at {:?}",
+        "COLD HEAD is at {:#?}",
         store.get_store(Temperature::Cold).get_ser::<Tip>(DBCol::BlockMisc, HEAD_KEY)
     );
 }

tools/state-viewer/Cargo.toml

@@ -47,3 +47,4 @@ nightly = [
 ]
 nightly_protocol = ["nearcore/nightly_protocol"]
 protocol_feature_flat_state = ["nearcore/protocol_feature_flat_state"]
+cold_store = ["near-store/cold_store"]