Fix: Add --dry-run to the tool for applying state-parts (#8739)

Also rework the state parts command API to be a single command with two subcommands instead of two separate commands.
near · Mar 20, 2023 · 0352621 · 0352621
1 parent 659199c
commit 0352621
Show file tree

Hide file tree

Showing 4 changed files with 210 additions and 144 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,10 +6,11 @@
 
 ### Non-protocol Changes
 
+* State-viewer tool to dump and apply state changes from/to a range of blocks. [#8628](https://github.com/near/nearcore/pull/8628)
 * Experimental option to dump state of every epoch to external storage. [#8661](https://github.com/near/nearcore/pull/8661)
-* State-viewer tool to dump and apply state changes from/to a range of blocks [#8628](https://github.com/near/nearcore/pull/8628)
-* Node can restart if State Sync gets interrupted [#8732](https://github.com/near/nearcore/pull/8732)
-* Add prometheus metrics for tracked shards, block height within epoch, if is block/chunk producer
+* Add prometheus metrics for tracked shards, block height within epoch, if is block/chunk producer. [#8728](https://github.com/near/nearcore/pull/8728)
+* Node can restart if State Sync gets interrupted. [#8732](https://github.com/near/nearcore/pull/8732)
+* Merged two `neard view-state` commands: `apply-state-parts` and `dump-state-parts` into a single `state-parts` command. [#8739](https://github.com/near/nearcore/pull/8739)
 
 ## 1.32.0
 

diff --git a/tools/state-viewer/src/apply_chain_range.rs b/tools/state-viewer/src/apply_chain_range.rs
@@ -1,10 +1,3 @@
-use std::fs::File;
-use std::io::Write;
-use std::sync::atomic::{AtomicU64, Ordering};
-use std::sync::Mutex;
-
-use rayon::iter::{IntoParallelIterator, ParallelIterator};
-
 use near_chain::chain::collect_receipts_from_response;
 use near_chain::migrations::check_if_block_is_first_with_chunk_of_version;
 use near_chain::types::ApplyTransactionResult;
@@ -19,6 +12,11 @@ use near_primitives::types::chunk_extra::ChunkExtra;
 use near_primitives::types::{BlockHeight, ShardId};
 use near_store::{get, DBCol, Store};
 use nearcore::NightshadeRuntime;
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+use std::fs::File;
+use std::io::Write;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::Mutex;
 
 fn timestamp_ms() -> u64 {
     use std::time::{SystemTime, UNIX_EPOCH};

diff --git a/tools/state-viewer/src/cli.rs b/tools/state-viewer/src/cli.rs
@@ -1,7 +1,6 @@
 use crate::commands::*;
 use crate::contract_accounts::ContractAccountFilter;
 use crate::rocksdb_stats::get_rocksdb_stats;
-use crate::state_parts::{apply_state_parts, dump_state_parts};
 use near_chain_configs::{GenesisChangeConfig, GenesisValidationMode};
 use near_primitives::account::id::AccountId;
 use near_primitives::hash::CryptoHash;
@@ -27,8 +26,6 @@ pub enum StateViewerSubCommand {
     /// even if it's not included in any block on disk
     #[clap(alias = "apply_receipt")]
     ApplyReceipt(ApplyReceiptCmd),
-    /// Apply all or a single state part of a shard.
-    ApplyStateParts(ApplyStatePartsCmd),
     /// Apply a transaction if it occurs in some chunk we know about,
     /// even if it's not included in any block on disk
     #[clap(alias = "apply_tx")]
@@ -52,8 +49,6 @@ pub enum StateViewerSubCommand {
     /// Generate a genesis file from the current state of the DB.
     #[clap(alias = "dump_state")]
     DumpState(DumpStateCmd),
-    /// Dump all or a single state part of a shard.
-    DumpStateParts(DumpStatePartsCmd),
     /// Writes state to a remote redis server.
     #[clap(alias = "dump_state_redis")]
     DumpStateRedis(DumpStateRedisCmd),
@@ -78,6 +73,8 @@ pub enum StateViewerSubCommand {
     /// Dumps or applies StateChanges.
     /// Experimental tool for shard shadowing development.
     StateChanges(StateChangesCmd),
+    /// Dump or apply state parts.
+    StateParts(StatePartsCmd),
     /// View head of the storage.
     #[clap(alias = "view_chain")]
     ViewChain(ViewChainCmd),
@@ -117,7 +114,6 @@ impl StateViewerSubCommand {
             StateViewerSubCommand::ApplyChunk(cmd) => cmd.run(home_dir, near_config, store),
             StateViewerSubCommand::ApplyRange(cmd) => cmd.run(home_dir, near_config, store),
             StateViewerSubCommand::ApplyReceipt(cmd) => cmd.run(home_dir, near_config, store),
-            StateViewerSubCommand::ApplyStateParts(cmd) => cmd.run(home_dir, near_config, store),
             StateViewerSubCommand::ApplyTx(cmd) => cmd.run(home_dir, near_config, store),
             StateViewerSubCommand::Chain(cmd) => cmd.run(home_dir, near_config, store),
             StateViewerSubCommand::CheckBlock => check_block_chunk_existence(near_config, store),
@@ -126,7 +122,6 @@ impl StateViewerSubCommand {
             StateViewerSubCommand::DumpAccountStorage(cmd) => cmd.run(home_dir, near_config, store),
             StateViewerSubCommand::DumpCode(cmd) => cmd.run(home_dir, near_config, store),
             StateViewerSubCommand::DumpState(cmd) => cmd.run(home_dir, near_config, store),
-            StateViewerSubCommand::DumpStateParts(cmd) => cmd.run(home_dir, near_config, store),
             StateViewerSubCommand::DumpStateRedis(cmd) => cmd.run(home_dir, near_config, store),
             StateViewerSubCommand::DumpTx(cmd) => cmd.run(home_dir, near_config, store),
             StateViewerSubCommand::EpochInfo(cmd) => cmd.run(home_dir, near_config, store),
@@ -136,6 +131,7 @@ impl StateViewerSubCommand {
             StateViewerSubCommand::RocksDBStats(cmd) => cmd.run(store_opener.path()),
             StateViewerSubCommand::State => state(home_dir, near_config, store),
             StateViewerSubCommand::StateChanges(cmd) => cmd.run(home_dir, near_config, store),
+            StateViewerSubCommand::StateParts(cmd) => cmd.run(home_dir, near_config, store),
             StateViewerSubCommand::ViewChain(cmd) => cmd.run(near_config, store),
             StateViewerSubCommand::ViewTrie(cmd) => cmd.run(store),
         }
@@ -219,42 +215,6 @@ impl ApplyReceiptCmd {
     }
 }
 
-#[derive(clap::Parser)]
-pub struct ApplyStatePartsCmd {
-    /// Selects an epoch. The dump will be of the state at the beginning of this epoch.
-    #[clap(subcommand)]
-    epoch_selection: crate::state_parts::EpochSelection,
-    /// Shard id.
-    #[clap(long)]
-    shard_id: ShardId,
-    /// State part id. Leave empty to go through every part in the shard.
-    #[clap(long)]
-    part_id: Option<u64>,
-    /// Where to write the state parts to.
-    #[clap(long)]
-    root_dir: Option<PathBuf>,
-    /// Store state parts in an S3 bucket.
-    #[clap(long)]
-    s3_bucket: Option<String>,
-    /// Store state parts in an S3 bucket.
-    #[clap(long)]
-    s3_region: Option<String>,
-}
-
-impl ApplyStatePartsCmd {
-    pub fn run(self, home_dir: &Path, near_config: NearConfig, store: Store) {
-        apply_state_parts(
-            self.epoch_selection,
-            self.shard_id,
-            self.part_id,
-            home_dir,
-            near_config,
-            store,
-            crate::state_parts::Location::new(self.root_dir, (self.s3_bucket, self.s3_region)),
-        );
-    }
-}
-
 #[derive(clap::Parser)]
 pub struct ApplyTxCmd {
     #[clap(long)]
@@ -402,49 +362,6 @@ impl DumpStateCmd {
     }
 }
 
-#[derive(clap::Parser)]
-pub struct DumpStatePartsCmd {
-    /// Selects an epoch. The dump will be of the state at the beginning of this epoch.
-    #[clap(subcommand)]
-    epoch_selection: crate::state_parts::EpochSelection,
-    /// Shard id.
-    #[clap(long)]
-    shard_id: ShardId,
-    /// Dump a single part id.
-    #[clap(long)]
-    part_id: Option<u64>,
-    /// Dump part ids starting from this part.
-    #[clap(long)]
-    part_from: Option<u64>,
-    /// Dump part ids up to this part (exclusive).
-    #[clap(long)]
-    part_to: Option<u64>,
-    /// Where to write the state parts to.
-    #[clap(long)]
-    root_dir: Option<PathBuf>,
-    /// Store state parts in an S3 bucket.
-    #[clap(long)]
-    s3_bucket: Option<String>,
-    /// S3 region to store state parts.
-    #[clap(long)]
-    s3_region: Option<String>,
-}
-
-impl DumpStatePartsCmd {
-    pub fn run(self, home_dir: &Path, near_config: NearConfig, store: Store) {
-        dump_state_parts(
-            self.epoch_selection,
-            self.shard_id,
-            self.part_from.or(self.part_id),
-            self.part_to.or(self.part_id.map(|x| x + 1)),
-            home_dir,
-            near_config,
-            store,
-            crate::state_parts::Location::new(self.root_dir, (self.s3_bucket, self.s3_region)),
-        );
-    }
-}
-
 #[derive(clap::Parser)]
 pub struct DumpStateRedisCmd {
     /// Optionally, can specify at which height to dump state.
@@ -655,3 +572,36 @@ impl ViewTrieCmd {
         }
     }
 }
+
+#[derive(clap::Parser)]
+pub struct StatePartsCmd {
+    /// Shard id.
+    #[clap(long)]
+    shard_id: ShardId,
+    /// Location of serialized state parts.
+    #[clap(long)]
+    root_dir: Option<PathBuf>,
+    /// Store state parts in an S3 bucket.
+    #[clap(long)]
+    s3_bucket: Option<String>,
+    /// Store state parts in an S3 bucket.
+    #[clap(long)]
+    s3_region: Option<String>,
+    /// Dump or Apply state parts.
+    #[clap(subcommand)]
+    command: crate::state_parts::StatePartsSubCommand,
+}
+
+impl StatePartsCmd {
+    pub fn run(self, home_dir: &Path, near_config: NearConfig, store: Store) {
+        self.command.run(
+            self.shard_id,
+            self.root_dir,
+            self.s3_bucket,
+            self.s3_region,
+            home_dir,
+            near_config,
+            store,
+        );
+    }
+}