Safe DB migrations using RocksDB Checkpoints

core/store/src/db.rs

@@ -16,6 +16,7 @@ use near_primitives::version::DbVersion;
 
 use crate::db::refcount::merge_refcounted_records;
 
+use rocksdb::checkpoint::Checkpoint;
 use std::path::Path;
 use std::sync::atomic::Ordering;
 
@@ -872,6 +873,10 @@ impl RocksDB {
             Ok(())
         }
     }
+
+    pub fn checkpoint(&self) -> Result<Checkpoint, DBError> {
+        Checkpoint::new(&self.db).map_err(|err| DBError(err))
+    }
 }
 
 fn available_space<P: AsRef<Path> + std::fmt::Debug>(

nearcore/Cargo.toml

@@ -8,33 +8,34 @@ rust-version = "1.56.0"
 edition = "2021"
 
 [dependencies]
-anyhow = "1.0.51"
-awc = "3.0.0-beta.5"
 actix = "=0.11.0-beta.2" # Pinned the version to avoid compilation errors
-actix_derive = "=0.6.0-beta.1" # Pinned dependency in addition to actix dependecy (remove this line once the pinning is not needed)
-actix-web = "=4.0.0-beta.6"
 actix-rt = "2"
+actix-web = "=4.0.0-beta.6"
+actix_derive = "=0.6.0-beta.1" # Pinned dependency in addition to actix dependecy (remove this line once the pinning is not needed)
+anyhow = "1.0.51"
+awc = "3.0.0-beta.5"
+borsh = "0.9"
 byteorder = "1.2"
-easy-ext = "0.2"
 chrono = { version = "0.4.4", features = ["serde"] }
+dirs = "3"
+easy-ext = "0.2"
 futures = "0.3"
 hyper = { version = "0.14", features = ["full"] }
 hyper-tls = "0.5.0"
+indicatif = "0.15.0"
+lazy-static-include = "3"
+near-rust-allocator-proxy = { version = "0.4", optional = true }
+num-rational = { version = "0.3", features = ["serde"] }
 rand = "0.7"
 rayon = "1.5"
+rocksdb = { version = "0.18.0", default-features = false, features = ["snappy", "lz4", "zstd", "zlib"] }
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
-dirs = "3"
-borsh = "0.9"
+smart-default = "0.6"
+tempfile = "3"
 thiserror = "1.0"
 tokio = { version = "1.1", features = ["fs"] }
 tracing = "0.1.13"
-smart-default = "0.6"
-num-rational = { version = "0.3", features = ["serde"] }
-near-rust-allocator-proxy = { version = "0.4", optional = true }
-lazy-static-include = "3"
-tempfile = "3"
-indicatif = "0.15.0"
 xz2 = "0.1.6"
 
 near-crypto = { path = "../core/crypto" }

nearcore/src/config.rs

@@ -322,6 +322,10 @@ fn default_trie_viewer_state_size_limit() -> Option<u64> {
     Some(50_000)
 }
 
+fn default_use_checkpoints_for_db_migration() -> bool {
+    true
+}
+
 #[derive(Serialize, Deserialize, Clone, Debug)]
 pub struct Consensus {
     /// Minimum number of peers to start syncing.
@@ -435,6 +439,12 @@ pub struct Config {
     /// If set, overrides value in genesis configuration.
     #[serde(skip_serializing_if = "Option::is_none")]
     pub max_gas_burnt_view: Option<Gas>,
+    /// Checkpoints let the user recover from interrupted DB migrations.
+    #[serde(default = "default_use_checkpoints_for_db_migration")]
+    pub use_checkpoints_for_db_migration: bool,
+    /// Absolute path to the root directory for DB checkpoints.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub db_checkpoints_path: Option<PathBuf>,
 }
 
 impl Default for Config {
@@ -461,6 +471,8 @@ impl Default for Config {
             view_client_throttle_period: default_view_client_throttle_period(),
             trie_viewer_state_size_limit: default_trie_viewer_state_size_limit(),
             max_gas_burnt_view: None,
+            db_checkpoints_path: None,
+            use_checkpoints_for_db_migration: true,
         }
     }
 }

nearcore/src/lib.rs

@@ -1,26 +1,27 @@
-use std::fs;
-use std::path::{Path, PathBuf};
-use std::sync::Arc;
-
+pub use crate::config::{init_configs, load_config, load_test_config, NearConfig, NEAR_BASE};
+use crate::migrations::{
+    migrate_12_to_13, migrate_18_to_19, migrate_19_to_20, migrate_22_to_23, migrate_23_to_24,
+    migrate_24_to_25, migrate_30_to_31,
+};
+pub use crate::runtime::NightshadeRuntime;
+pub use crate::shard_tracker::TrackedConfig;
 use actix::{Actor, Addr, Arbiter};
 use actix_rt::ArbiterHandle;
 use actix_web;
 use anyhow::Context;
-#[cfg(feature = "performance_stats")]
-use near_rust_allocator_proxy::reset_memory_usage_max;
-use tracing::{error, info, trace};
-
 use near_chain::ChainGenesis;
 #[cfg(feature = "test_features")]
 use near_client::AdversarialControls;
 use near_client::{start_client, start_view_client, ClientActor, ViewClientActor};
-
 use near_network::routing::start_routing_table_actor;
 use near_network::test_utils::NetworkRecipient;
 use near_network::PeerManagerActor;
 use near_primitives::network::PeerId;
 #[cfg(feature = "rosetta_rpc")]
 use near_rosetta_rpc::start_rosetta_rpc;
+#[cfg(feature = "performance_stats")]
+use near_rust_allocator_proxy::reset_memory_usage_max;
+use near_store::db::RocksDB;
 use near_store::migrations::{
     fill_col_outcomes_by_hash, fill_col_transaction_refcount, get_store_version, migrate_10_to_11,
     migrate_11_to_12, migrate_13_to_14, migrate_14_to_15, migrate_17_to_18, migrate_20_to_21,
@@ -29,14 +30,10 @@ use near_store::migrations::{
 };
 use near_store::{create_store, Store};
 use near_telemetry::TelemetryActor;
-
-pub use crate::config::{init_configs, load_config, load_test_config, NearConfig, NEAR_BASE};
-use crate::migrations::{
-    migrate_12_to_13, migrate_18_to_19, migrate_19_to_20, migrate_22_to_23, migrate_23_to_24,
-    migrate_24_to_25, migrate_30_to_31,
-};
-pub use crate::runtime::NightshadeRuntime;
-pub use crate::shard_tracker::TrackedConfig;
+use std::path::{Path, PathBuf};
+use std::sync::Arc;
+use std::{fs, io};
+use tracing::{error, info, trace, warn};
 
 pub mod append_only_map;
 pub mod config;
@@ -74,17 +71,87 @@ pub fn get_default_home() -> PathBuf {
     PathBuf::default()
 }
 
+fn db_checkpoint_root_dir<'a>(path: &'a Path, near_config: &'a NearConfig) -> &'a Path {
+    if let Some(db_checkpoints_path) = &near_config.config.db_checkpoints_path {
+        &db_checkpoints_path
+    } else {
+        path
+    }
+}
+
+const DB_CHECKPOINT_PREFIX: &str = "db_checkpoint_";
+
+fn find_db_checkpoint(path: &Path, near_config: &NearConfig) -> Result<Option<PathBuf>, io::Error> {
+    let db_checkpoints_path = db_checkpoint_root_dir(path, near_config);
+    for entry in std::fs::read_dir(db_checkpoints_path)?.filter(|entry| {
+        if let Ok(entry) = entry {
+            if let Some(file_name) = entry.file_name().to_str() {
+                return file_name.starts_with(DB_CHECKPOINT_PREFIX);
+            }
+        }
+        return false;
+    }) {
+        return Ok(Some(PathBuf::from(entry?.path())));
+    }
+    Ok(None)
+}
+
+fn create_db_checkpoint(path: &Path, near_config: &NearConfig) -> Result<PathBuf, anyhow::Error> {
+    let checkpoint_dir =
+        format!("{}{}", DB_CHECKPOINT_PREFIX, &chrono::Utc::now().format("%Y%m%d%H%M%S"));
+    let checkpoint_path = db_checkpoint_root_dir(path, near_config);
+    let checkpoint_path = checkpoint_path.join(Path::new(&checkpoint_dir));
+
+    let db = RocksDB::new(path)?;
+    let checkpoint = db.checkpoint()?;
+    info!(target:"near","Creating a DB snapshot in '{}'",checkpoint_path.display());
+    checkpoint.create_checkpoint(&checkpoint_path)?;
+
+    Ok(checkpoint_path)
+}
+
 /// Function checks current version of the database and applies migrations to the database.
 pub fn apply_store_migrations(path: &Path, near_config: &NearConfig) {
     let db_version = get_store_version(path);
     if db_version > near_primitives::version::DB_VERSION {
         error!(target: "near", "DB version {} is created by a newer version of neard, please update neard or delete data", db_version);
         std::process::exit(1);
     }
+
     if db_version == near_primitives::version::DB_VERSION {
         return;
     }
 
+    if near_config.config.use_checkpoints_for_db_migration {
+        let existing_checkpoint_path = find_db_checkpoint(path, near_config);
+        match existing_checkpoint_path {
+            Ok(existing_checkpoint_path) => match existing_checkpoint_path {
+                Some(existing_checkpoint_path) => {
+                    panic!("Detected an existing DB migration checkpoint: '{}'. Probably a DB migration got interrupted and your DB is corrupted. Please replace the contents of '{}' with data from that checkpoint, delete the checkpoint and try again.",existing_checkpoint_path.display(), path.display());
+                }
+                None => {}
+            },
+            Err(err) => {
+                warn!(target: "near", "Failed to check if a DB checkpoint already exists: {:#?}", err);
+            }
+        }
+    }
+
+    let checkpoint_path = if near_config.config.use_checkpoints_for_db_migration {
+        match create_db_checkpoint(path, near_config) {
+            Ok(checkpoint_path) => {
+                info!(target: "near", "Created a DB checkpoint before a DB migration: '{}'. Please recover from this checkpoint if the migration gets interrupted.", checkpoint_path.display());
+                Some(checkpoint_path)
+            }
+            Err(err) => {
+                warn!(target: "near", "Failed to create a DB checkpoint before a DB migration: {}", err);
+                None
+            }
+        }
+    } else {
+        None
+    };
+
     // Add migrations here based on `db_version`.
     if db_version <= 1 {
         // version 1 => 2: add gc column
@@ -274,6 +341,20 @@ pub fn apply_store_migrations(path: &Path, near_config: &NearConfig) {
         let db_version = get_store_version(path);
         debug_assert_eq!(db_version, near_primitives::version::DB_VERSION);
     }
+
+    if near_config.config.use_checkpoints_for_db_migration {
+        if let Some(checkpoint_path) = checkpoint_path {
+            info!(target: "near", "Deleting DB checkpoint at '{}'", checkpoint_path.display());
+            match std::fs::remove_dir_all(&checkpoint_path) {
+                Ok(_) => {
+                    info!(target: "near","Deleted DB checkpoint at '{}'", checkpoint_path.display());
+                }
+                Err(err) => {
+                    error!(target: "near","Failed to delete a DB checkpoint at '{}'. Error: {:#?}. Please delete it manually before the next start of the node.", checkpoint_path.display(), err);
+                }
+            }
+        }
+    }
 }
 
 pub fn init_and_migrate_store(home_dir: &Path, near_config: &NearConfig) -> Store {