fix(state-snapshot): Tool to make DB snapshots (#9308)

Co-authored-by: near-bulldozer[bot] <73298989+near-bulldozer[bot]@users.noreply.github.com>

core/store/src/opener.rs

@@ -604,7 +604,7 @@ pub fn checkpoint_hot_storage_and_cleanup_columns(
     config.path = Some(checkpoint_path);
     let archive = hot_store.get_db_kind()? == Some(DbKind::Archive);
     let opener = StoreOpener::new(checkpoint_base_path, archive, &config, None);
-    let node_storage = opener.open()?;
+    let node_storage = opener.open_in_mode(Mode::ReadWriteExisting)?;
 
     if let Some(columns_to_keep) = columns_to_keep {
         let columns_to_keep_set: std::collections::HashSet<DBCol> =

tools/database/Cargo.toml

@@ -13,6 +13,7 @@ anyhow.workspace = true
 clap.workspace = true
 rayon.workspace = true
 strum.workspace = true
+tempfile.workspace = true
 
 nearcore.workspace = true
 near-store.workspace = true

tools/database/README.md

@@ -11,7 +11,7 @@ of keys and values within RocksDB.
 
 To run the script, use the following example:
 ```bash
-cargo run --bin neard -- --home /home/ubuntu/.nerd database analyse-data-size-distribution --column State --top_k 50
+cargo run --bin neard -- --home /home/ubuntu/.near database analyse-data-size-distribution --column State --top_k 50
 ```
 The arguments are as follows:
 
@@ -45,7 +45,10 @@ It is intended as a collection of commands that perform small db modifications.
 ### change-db-kind
 Changes DbKind of a DB described in config (cold or hot).  
 Example usage:
-`neard database change-db-kind --new-kind RPC change-cold`  
+```bash
+cargo run --bin neard -- --home /home/ubuntu/.near database change-db-kind --new-kind RPC change-cold
+```
+
 In this example we change DbKind of the cold db to RPC (for some reason).  
 Notice, that you cannot perform this exact command twice in a row,
 because you will not be able to open cold db in the first place.  
@@ -59,3 +62,19 @@ Then you can call
 `neard database change-db-kind --new-kind Cold change-hot`.
 Notice that even though in your mind this db is cold, in your config this db hot, so you have to pass `change-hot`.
 
+## Make a DB Snapshot
+
+Makes a copy of a DB (hot store only) at a specified location. If the
+destination is within the same filesystem, the copy will be made instantly and
+take no additional disk space due to hardlinking all the files.
+
+Example usage:
+```bash
+cargo run --bin neard -- --home /home/ubuntu/.near database make_snapshot --destination /home/ubuntu/.near/data/snapshot
+```
+
+In this example all `.sst` files from `/home/ubuntu/.near/data` will be also
+available in `/home/ubuntu/.near/data/snapshot`
+
+This command can be helpful before attempting activities that can potentially
+corrupt the database.

tools/database/src/commands.rs

@@ -1,5 +1,6 @@
 use crate::adjust_database::ChangeDbKindCommand;
 use crate::analyse_data_size_distribution::AnalyseDataSizeDistributionCommand;
+use crate::make_snapshot::MakeSnapshotCommand;
 use clap::Parser;
 use std::path::PathBuf;
 
@@ -17,6 +18,9 @@ enum SubCommand {
 
     /// Change DbKind of hot or cold db.
     ChangeDbKind(ChangeDbKindCommand),
+
+    /// Make snapshot of the database
+    MakeSnapshot(MakeSnapshotCommand),
 }
 
 impl DatabaseCommand {
@@ -31,6 +35,14 @@ impl DatabaseCommand {
                 .unwrap_or_else(|e| panic!("Error loading config: {:#}", e));
                 cmd.run(home, &near_config)
             }
+            SubCommand::MakeSnapshot(cmd) => {
+                let near_config = nearcore::config::load_config(
+                    &home,
+                    near_chain_configs::GenesisValidationMode::UnsafeFast,
+                )
+                .unwrap_or_else(|e| panic!("Error loading config: {:#}", e));
+                cmd.run(home, near_config.config.archive, &near_config.config.store)
+            }
         }
     }
 }

tools/database/src/lib.rs

@@ -1,3 +1,4 @@
 mod adjust_database;
 mod analyse_data_size_distribution;
 pub mod commands;
+mod make_snapshot;

tools/database/src/make_snapshot.rs

@@ -0,0 +1,83 @@
+use near_store::{checkpoint_hot_storage_and_cleanup_columns, Mode, NodeStorage, StoreConfig};
+use std::path::{Path, PathBuf};
+
+#[derive(clap::Args)]
+pub(crate) struct MakeSnapshotCommand {
+    /// Destination directory.
+    #[clap(long)]
+    destination: PathBuf,
+}
+
+impl MakeSnapshotCommand {
+    pub(crate) fn run(
+        &self,
+        home_dir: &Path,
+        archive: bool,
+        store_config: &StoreConfig,
+    ) -> anyhow::Result<()> {
+        let opener = NodeStorage::opener(home_dir, archive, store_config, None);
+        let node_storage = opener.open_in_mode(Mode::ReadWriteExisting)?;
+        checkpoint_hot_storage_and_cleanup_columns(
+            &node_storage.get_hot_store(),
+            &self.destination,
+            None,
+        )?;
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::make_snapshot::MakeSnapshotCommand;
+    use near_store::{DBCol, Mode, NodeStorage, StoreConfig};
+
+    /// Populates a DB, makes a checkpoint, makes changes to the DB.
+    /// Checks that the checkpoint DB can be opened and doesn't contain the latest changes.
+    #[test]
+    fn test() {
+        let home_dir = tempfile::tempdir().unwrap();
+        let store_config = StoreConfig::test_config();
+        let opener = NodeStorage::opener(home_dir.path(), false, &store_config, None);
+
+        let keys = vec![vec![0], vec![1], vec![2], vec![3]];
+
+        {
+            // Populate the DB.
+            let node_storage = opener.open().unwrap();
+            let mut store_update = node_storage.get_hot_store().store_update();
+            for key in &keys {
+                store_update.insert(DBCol::Block, key, &vec![42]);
+            }
+            store_update.commit().unwrap();
+            println!("Populated");
+            // Drops node_storage, which unlocks the DB.
+        }
+
+        let destination = home_dir.path().join("data").join("snapshot");
+        let cmd = MakeSnapshotCommand { destination: destination.clone() };
+        cmd.run(home_dir.path(), false, &store_config).unwrap();
+        println!("Made a checkpoint");
+
+        {
+            // Make a change to the original DB.
+            let node_storage = opener.open().unwrap();
+            let mut store_update = node_storage.get_hot_store().store_update();
+            store_update.delete_all(DBCol::Block);
+            store_update.commit().unwrap();
+            println!("Deleted");
+        }
+
+        let node_storage = opener.open_in_mode(Mode::ReadOnly).unwrap();
+        let snapshot_node_storage = NodeStorage::opener(&destination, false, &store_config, None)
+            .open_in_mode(Mode::ReadOnly)
+            .unwrap();
+        for key in keys {
+            let exists_original = node_storage.get_hot_store().exists(DBCol::Block, &key).unwrap();
+            let exists_snapshot =
+                snapshot_node_storage.get_hot_store().exists(DBCol::Block, &key).unwrap();
+            println!("{exists_original},{exists_snapshot},{key:?}");
+            assert!(!exists_original);
+            assert!(exists_snapshot);
+        }
+    }
+}