From cf3d213b3f0c27a4f5a27080dde914299ce94101 Mon Sep 17 00:00:00 2001 From: Sayan Nandan Date: Thu, 29 Oct 2020 12:46:38 +0530 Subject: [PATCH] Enable data to be restored from a snapshot The user can now run `tdb -r ` to restore data from the snapshot. Also, we'll show a note in the logs when trying to restore from a snapshot Signed-off-by: Sayan Nandan --- examples/config-files/snapshot.toml | 2 +- server/src/cli.yml | 7 +++++++ server/src/config/mod.rs | 18 ++++++++++++------ server/src/coredb/mod.rs | 9 +++++++-- server/src/dbnet.rs | 4 +++- server/src/diskstore/mod.rs | 6 +++++- server/src/main.rs | 23 +++++++++++++++++------ server/src/tests/mod.rs | 2 +- tdb-macros/src/lib.rs | 3 ++- 9 files changed, 55 insertions(+), 19 deletions(-) diff --git a/examples/config-files/snapshot.toml b/examples/config-files/snapshot.toml index 8e51addc..394368c3 100644 --- a/examples/config-files/snapshot.toml +++ b/examples/config-files/snapshot.toml @@ -13,6 +13,6 @@ every = 120 [snapshot] enabled = true # set to false to disable snapshots # Create a snapshot every hour (1 hour = 60 minutes = 60 * 60 seconds = 3600 seconds) -every = 3600 +every = 10 # How many of the snapshots to keep atmost = 4 # keep the four most recent snapshots diff --git a/server/src/cli.yml b/server/src/cli.yml index 0f1dfb52..f0c842fd 100644 --- a/server/src/cli.yml +++ b/server/src/cli.yml @@ -31,3 +31,10 @@ args: value_name: cfgfile help: Use a configuration file to start tdb takes_value: true + - restore: + short: r + required: false + long: restore + value_name: restorefile + help: Restore data from a previous snapshot + takes_value: true diff --git a/server/src/config/mod.rs b/server/src/config/mod.rs index 5fc01f6d..ef9ea245 100644 --- a/server/src/config/mod.rs +++ b/server/src/config/mod.rs @@ -30,6 +30,7 @@ use std::fs; #[cfg(test)] use std::net::Ipv6Addr; use std::net::{IpAddr, Ipv4Addr}; +use std::path::PathBuf; use tokio::net::ToSocketAddrs; use toml; @@ -364,9 +365,9 @@ use clap::{load_yaml, App}; /// The type of configuration: /// - We either used a custom configuration file given to us by the user (`Custom`) OR /// - We used the default configuration (`Def`) -pub enum ConfigType { - Def(T), - Custom(T), +pub enum ConfigType { + Def(T, Option), + Custom(T, Option), } #[derive(Debug)] @@ -396,9 +397,14 @@ impl fmt::Display for ConfigError { /// This parses a configuration file if it is supplied as a command line argument /// or it returns the default configuration. **If** the configuration file /// contains an error, then this returns it as an `Err` variant -pub fn get_config_file_or_return_cfg() -> Result, ConfigError> { +pub fn get_config_file_or_return_cfg() -> Result, ConfigError> { let cfg_layout = load_yaml!("../cli.yml"); let matches = App::from_yaml(cfg_layout).get_matches(); + let restorefile = matches.value_of("restore").map(|val| { + let mut path = PathBuf::from("snapshots/"); + path.push(val); + path + }); let filename = matches.value_of("config"); if let Some(filename) = filename { match ParsedConfig::new_from_file(filename.to_owned()) { @@ -420,12 +426,12 @@ pub fn get_config_file_or_return_cfg() -> Result, Confi )); } } - return Ok(ConfigType::Custom(cfg)); + return Ok(ConfigType::Custom(cfg, restorefile)); } Err(e) => return Err(e), } } else { - Ok(ConfigType::Def(ParsedConfig::default())) + Ok(ConfigType::Def(ParsedConfig::default(), restorefile)) } } diff --git a/server/src/coredb/mod.rs b/server/src/coredb/mod.rs index a1d4c51f..fd013dd4 100644 --- a/server/src/coredb/mod.rs +++ b/server/src/coredb/mod.rs @@ -35,6 +35,7 @@ use parking_lot::RwLock; use parking_lot::RwLockReadGuard; use parking_lot::RwLockWriteGuard; use std::collections::HashMap; +use std::path::PathBuf; use std::sync::Arc; use tokio; use tokio::sync::Notify; @@ -238,8 +239,12 @@ impl CoreDB { /// /// This also checks if a local backup of previously saved data is available. /// If it is - it restores the data. Otherwise it creates a new in-memory table - pub fn new(bgsave: BGSave, snapshot_cfg: SnapshotConfig) -> TResult { - let coretable = diskstore::get_saved(Some(PERSIST_FILE.to_path_buf()))?; + pub fn new( + bgsave: BGSave, + snapshot_cfg: SnapshotConfig, + restore_file: Option, + ) -> TResult { + let coretable = diskstore::get_saved(restore_file)?; let mut background_tasks: usize = 0; if !bgsave.is_disabled() { background_tasks += 1; diff --git a/server/src/dbnet.rs b/server/src/dbnet.rs index 03e1be5c..1b50bcde 100644 --- a/server/src/dbnet.rs +++ b/server/src/dbnet.rs @@ -41,6 +41,7 @@ use crate::CoreDB; use libtdb::util::terminal; use libtdb::TResult; use std::future::Future; +use std::path::PathBuf; use std::process; use std::sync::Arc; use tokio::net::TcpListener; @@ -189,10 +190,11 @@ pub async fn run( bgsave_cfg: BGSave, snapshot_cfg: SnapshotConfig, sig: impl Future, + restore_filepath: Option, ) { let (signal, _) = broadcast::channel(1); let (terminate_tx, terminate_rx) = mpsc::channel(1); - let db = match CoreDB::new(bgsave_cfg, snapshot_cfg) { + let db = match CoreDB::new(bgsave_cfg, snapshot_cfg, restore_filepath) { Ok(d) => d, Err(e) => { eprintln!("ERROR: {}", e); diff --git a/server/src/diskstore/mod.rs b/server/src/diskstore/mod.rs index 2ea444d1..4ffecf8d 100644 --- a/server/src/diskstore/mod.rs +++ b/server/src/diskstore/mod.rs @@ -44,7 +44,11 @@ lazy_static::lazy_static! { /// otherwise the `data.bin` file is deserialized and parsed into a `HashMap` pub fn get_saved(location: Option) -> TResult>> { let file = match fs::read(if let Some(loc) = location { - PathBuf::from(loc) + log::info!( + "Attempting to restore data from file: '{}'", + loc.to_string_lossy() + ); + loc } else { PERSIST_FILE.to_path_buf() }) { diff --git a/server/src/main.rs b/server/src/main.rs index b0e13c9e..dba10978 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -74,22 +74,29 @@ async fn main() { .init(); // Start the server which asynchronously waits for a CTRL+C signal // which will safely shut down the server - let (tcplistener, bgsave_config, snapshot_config) = check_args_or_connect().await; + let (tcplistener, bgsave_config, snapshot_config, restore_filepath) = + check_args_or_connect().await; run( tcplistener, bgsave_config, snapshot_config, signal::ctrl_c(), + restore_filepath, ) .await; } /// This function checks the command line arguments and binds to an appropriate /// port and host, as per the supplied configuration options -async fn check_args_or_connect() -> (TcpListener, BGSave, SnapshotConfig) { +async fn check_args_or_connect() -> ( + TcpListener, + BGSave, + SnapshotConfig, + Option, +) { let cfg = config::get_config_file_or_return_cfg(); let binding_and_cfg = match cfg { - Ok(config::ConfigType::Custom(cfg)) => { + Ok(config::ConfigType::Custom(cfg, file)) => { if cfg.is_artful() { println!("{}\n{}", TEXT, MSG); } else { @@ -100,15 +107,17 @@ async fn check_args_or_connect() -> (TcpListener, BGSave, SnapshotConfig) { TcpListener::bind(cfg.get_host_port_tuple()).await, cfg.bgsave, cfg.snapshot, + file, ) } - Ok(config::ConfigType::Def(cfg)) => { + Ok(config::ConfigType::Def(cfg, file)) => { println!("{}\n{}", TEXT, MSG); log::warn!("No configuration file supplied. Using default settings"); ( TcpListener::bind(cfg.get_host_port_tuple()).await, cfg.bgsave, cfg.snapshot, + file, ) } Err(e) => { @@ -117,8 +126,10 @@ async fn check_args_or_connect() -> (TcpListener, BGSave, SnapshotConfig) { } }; match binding_and_cfg { - (Ok(b), bgsave_cfg, snapshot_cfg) => (b, bgsave_cfg, snapshot_cfg), - (Err(e), _, _) => { + (Ok(b), bgsave_cfg, snapshot_cfg, restore_file) => { + (b, bgsave_cfg, snapshot_cfg, restore_file) + } + (Err(e), _, _, _) => { log::error!("Failed to bind to socket with error: '{}'", e); std::process::exit(0x100); } diff --git a/server/src/tests/mod.rs b/server/src/tests/mod.rs index 1a685140..f732f7e3 100644 --- a/server/src/tests/mod.rs +++ b/server/src/tests/mod.rs @@ -48,7 +48,7 @@ async fn start_test_server(port: u16, db: Option) -> SocketAddr { let db = if let Some(db) = db { db } else { - CoreDB::new(BGSave::Disabled, SnapshotConfig::default()).unwrap() + CoreDB::new(BGSave::Disabled, SnapshotConfig::default(), None).unwrap() }; let listener = TcpListener::bind(socket) .await diff --git a/tdb-macros/src/lib.rs b/tdb-macros/src/lib.rs index c8df74da..61bb3d00 100644 --- a/tdb-macros/src/lib.rs +++ b/tdb-macros/src/lib.rs @@ -61,7 +61,8 @@ fn parse_dbtest(mut input: syn::ItemFn, rand: u16) -> Result