Change state dump to binary format (#2086)

Instead of dumping state in a json file which is slow, this PR changes it to dumping state in a binary format. Fixes #2070.

Test plan
---------
Manually test this with a local node to make sure that the state is preserved correctly after state dumping.

near/src/config.rs

@@ -137,6 +137,8 @@ pub const CONFIG_FILENAME: &str = "config.json";
 pub const GENESIS_CONFIG_FILENAME: &str = "genesis.json";
 pub const NODE_KEY_FILE: &str = "node_key.json";
 pub const VALIDATOR_KEY_FILE: &str = "validator_key.json";
+pub const STATE_FILE: &str = "state";
+pub const GENESIS_ROOTS_FILE: &str = "genesis_roots";
 
 pub const DEFAULT_TELEMETRY_URL: &str = "https://explorer.nearprotocol.com/api/nodes";
 
@@ -239,6 +241,8 @@ impl Default for Consensus {
 pub struct Config {
     pub genesis_file: String,
     pub validator_key_file: String,
+    pub state_file: String,
+    pub state_roots_file: String,
     pub node_key_file: String,
     pub rpc: RpcConfig,
     pub telemetry: TelemetryConfig,
@@ -253,6 +257,8 @@ impl Default for Config {
         Config {
             genesis_file: GENESIS_CONFIG_FILENAME.to_string(),
             validator_key_file: VALIDATOR_KEY_FILE.to_string(),
+            state_file: STATE_FILE.to_string(),
+            state_roots_file: GENESIS_ROOTS_FILE.to_string(),
             node_key_file: NODE_KEY_FILE.to_string(),
             rpc: RpcConfig::default(),
             telemetry: TelemetryConfig::default(),

near/src/runtime.rs

@@ -41,12 +41,10 @@ use node_runtime::adapter::ViewRuntimeAdapter;
 use node_runtime::state_viewer::TrieViewer;
 use node_runtime::{ApplyState, Runtime, StateRecord, ValidatorAccountsUpdate};
 
-use crate::config::GenesisConfig;
+use crate::config::{GenesisConfig, GENESIS_ROOTS_FILE, STATE_FILE};
 use crate::shard_tracker::{account_id_to_shard_id, ShardTracker};
 
 const POISONED_LOCK_ERR: &str = "The lock was poisoned.";
-const STATE_DUMP_FILE: &str = "state_dump";
-const GENESIS_ROOTS_FILE: &str = "genesis_roots";
 
 /// Defines Nightshade state transition, validator rotation and block weight for fork choice rule.
 /// TODO: this possibly should be merged with the runtime cargo or at least reconciled on the interfaces.
@@ -142,7 +140,7 @@ impl NightshadeRuntime {
     fn genesis_state_from_dump(&self) -> (StoreUpdate, Vec<StateRoot>) {
         let store_update = self.store.store_update();
         let mut state_file = self.home_dir.clone();
-        state_file.push(STATE_DUMP_FILE);
+        state_file.push(STATE_FILE);
         self.store
             .load_from_file(ColState, state_file.as_path())
             .expect("Failed to read state dump");
@@ -370,7 +368,7 @@ impl RuntimeAdapter for NightshadeRuntime {
         let has_records = !self.genesis_config.records.is_empty();
         let has_dump = {
             let mut state_dump = self.home_dir.clone();
-            state_dump.push(STATE_DUMP_FILE);
+            state_dump.push(STATE_FILE);
             state_dump.exists()
         };
         if has_dump {

test-utils/state-viewer/src/main.rs

@@ -1,8 +1,8 @@
 use std::convert::TryFrom;
-use std::path::Path;
+use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
-use borsh::BorshDeserialize;
+use borsh::{BorshDeserialize, BorshSerialize};
 use clap::{App, Arg, SubCommand};
 
 use ansi_term::Color::Red;
@@ -18,9 +18,11 @@ use near_primitives::test_utils::init_integration_logger;
 use near_primitives::types::{BlockHeight, StateRoot};
 use near_primitives::utils::{col, ACCOUNT_DATA_SEPARATOR};
 use near_store::test_utils::create_test_store;
-use near_store::{create_store, Store, TrieIterator};
+use near_store::{create_store, ColState, Store, TrieIterator};
 use node_runtime::StateRecord;
 use std::collections::HashMap;
+use std::fs::File;
+use std::io::Write;
 
 fn to_printable(blob: &[u8]) -> String {
     if blob.len() > 60 {
@@ -143,6 +145,33 @@ fn load_trie(
     (runtime, state_roots, last_block.header.inner_lite.height)
 }
 
+/// Dump state in binary format. Also write state roots to a separate file.
+fn dump_state(home_dir: &PathBuf, store: Arc<Store>, mut near_config: NearConfig) {
+    // Step 1: dumping state roots.
+    let mut chain_store = ChainStore::new(store.clone());
+    let head = chain_store.head().unwrap();
+    let last_block = chain_store.get_block(&head.last_block_hash).unwrap().clone();
+    let mut state_roots = vec![];
+    for chunk in last_block.chunks.iter() {
+        state_roots.push(chunk.inner.prev_state_root.clone());
+    }
+    println!("Saving state at {:?} @ {}", state_roots, head.height);
+    let mut roots_files = home_dir.clone();
+    roots_files.push("genesis_roots_dump");
+    let mut file = File::create(roots_files).unwrap();
+    let data = state_roots.try_to_vec().unwrap();
+    file.write_all(&data).unwrap();
+    // Step 2: dumping state.
+    let mut dump_path = home_dir.clone();
+    dump_path.push("state_dump");
+    store.save_to_file(ColState, dump_path.as_path()).unwrap();
+    // Step 3: dumping genesis config
+    let mut genesis_file = home_dir.clone();
+    genesis_file.push("genesis_config.json");
+    near_config.genesis_config.records = vec![];
+    near_config.genesis_config.write_to_file(&genesis_file);
+}
+
 pub fn format_hash(h: CryptoHash) -> String {
     to_base(&h)[..7].to_string()
 }
@@ -268,15 +297,7 @@ fn main() {
         )
         .subcommand(SubCommand::with_name("peers"))
         .subcommand(SubCommand::with_name("state"))
-        .subcommand(
-            SubCommand::with_name("dump_state").arg(
-                Arg::with_name("output")
-                    .long("output")
-                    .required(true)
-                    .help("Output path for new genesis given current blockchain state")
-                    .takes_value(true),
-            ),
-        )
+        .subcommand(SubCommand::with_name("dump_state"))
         .subcommand(
             SubCommand::with_name("chain")
                 .arg(
@@ -316,7 +337,7 @@ fn main() {
         .get_matches();
 
     let home_dir = matches.value_of("home").map(|dir| Path::new(dir)).unwrap();
-    let mut near_config = load_config(home_dir);
+    let near_config = load_config(home_dir);
 
     let store = create_store(&get_store_path(&home_dir));
 
@@ -338,26 +359,8 @@ fn main() {
                 }
             }
         }
-        ("dump_state", Some(args)) => {
-            let (runtime, state_roots, height) = load_trie(store, home_dir, &near_config);
-            let output_path = args.value_of("output").map(|path| Path::new(path)).unwrap();
-            println!(
-                "Saving state at {:?} @ {} into {}",
-                state_roots,
-                height,
-                output_path.display()
-            );
-            near_config.genesis_config.records = vec![];
-            for state_root in state_roots {
-                let trie = TrieIterator::new(&runtime.trie, &state_root).unwrap();
-                for item in trie {
-                    let (key, value) = item.unwrap();
-                    if let Some(sr) = kv_to_state_record(key, value) {
-                        near_config.genesis_config.records.push(sr);
-                    }
-                }
-            }
-            near_config.genesis_config.write_to_file(&output_path);
+        ("dump_state", _) => {
+            dump_state(&PathBuf::from(home_dir), store, near_config);
         }
         ("chain", Some(args)) => {
             let start_index =