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,11 @@ impl RocksDB {
             Ok(())
         }
     }
+
+    /// Creates a Checkpoint object that can be used to actually create a checkpoint on disk.
+    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,13 @@ 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.
+    /// If not set, defaults to the database location, i.e. '$home/data/'.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub db_checkpoints_path: Option<PathBuf>,
 }
 
 impl Default for Config {
@@ -461,6 +472,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::fs;
+use std::path::{Path, PathBuf};
+use std::sync::Arc;
+use tracing::{error, info, trace};
 
 pub mod append_only_map;
 pub mod config;
@@ -74,17 +71,89 @@ pub fn get_default_home() -> PathBuf {
     PathBuf::default()
 }
 
+/// Returns the path of a directory where DB checkpoint can be found or created.
+/// Default location is the same as the database location: `path`.
+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
+    }
+}
+
+/// Returns the path of the DB checkpoint.
+fn db_checkpoint_path<'a>(path: &'a Path, near_config: &'a NearConfig) -> PathBuf {
+    let checkpoint_path = db_checkpoint_root_dir(path, near_config);
+    let checkpoint_path = checkpoint_path.join(DB_CHECKPOINT_NAME);
+    checkpoint_path
+}
+
+const DB_CHECKPOINT_NAME: &str = "db_snapshot_before_migration";
+
+/// Returns the path of the first DB checkpoint found on disk.
+fn check_db_checkpoint_exists(path: &Path, near_config: &NearConfig) -> Option<PathBuf> {
+    let path = db_checkpoint_path(path, near_config);
+    if path.exists() {
+        Some(path)
+    } else {
+        None
+    }
+}
+
+/// Creates a consistent DB checkpoint and returns its path.
+/// By default it creates checkpoints in the DB directory, but can be overridden by the config.
+fn create_db_checkpoint(path: &Path, near_config: &NearConfig) -> Result<PathBuf, anyhow::Error> {
+    let checkpoint_path = db_checkpoint_path(path, near_config);
+
+    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)?;
+    info!(target:"near","Created a DB snapshot in '{}'",checkpoint_path.display());
+
+    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;
     }
 
+    // Before attemping a DB migration check if the current DB is corrupted. We assume that
+    // checkpoints get created before each migration attempt and get removed after successful
+    // attempts.
+    if near_config.config.use_checkpoints_for_db_migration {
+        match check_db_checkpoint_exists(path, near_config) {
+            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 => {}
+        }
+    }
+
+    // Before starting a DB migration, create a consistent snapshot of the database. If a migration
+    // fails, it can be used to quickly restore the database to its original state.
+    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) => {
+                panic!("Failed to create a DB checkpoint before a DB migration: {}. If this persists, you can disable checkpoints in `config.json`:\n```\"use_checkpoints_for_db_migration\": false\n```", err);
+            }
+        }
+    } else {
+        None
+    };
+
     // Add migrations here based on `db_version`.
     if db_version <= 1 {
         // version 1 => 2: add gc column
@@ -274,6 +343,21 @@ 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);
     }
+
+    // DB migration was successful, remove the checkpoint to avoid it taking up precious disk space.
+    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 {