Enable data to be restored from a snapshot

The user can now run `tdb -r <snapshotname>` 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 <nandansayan@outlook.com>

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

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

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<T> {
-    Def(T),
-    Custom(T),
+pub enum ConfigType<T, U> {
+    Def(T, Option<U>),
+    Custom(T, Option<U>),
 }
 
 #[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<ConfigType<ParsedConfig>, ConfigError> {
+pub fn get_config_file_or_return_cfg() -> Result<ConfigType<ParsedConfig, PathBuf>, 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<ConfigType<ParsedConfig>, 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))
     }
 }
 

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<Self> {
-        let coretable = diskstore::get_saved(Some(PERSIST_FILE.to_path_buf()))?;
+    pub fn new(
+        bgsave: BGSave,
+        snapshot_cfg: SnapshotConfig,
+        restore_file: Option<PathBuf>,
+    ) -> TResult<Self> {
+        let coretable = diskstore::get_saved(restore_file)?;
         let mut background_tasks: usize = 0;
         if !bgsave.is_disabled() {
             background_tasks += 1;

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<PathBuf>,
 ) {
     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);

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<PathBuf>) -> TResult<Option<HashMap<String, Data>>> {
     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()
     }) {

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<std::path::PathBuf>,
+) {
     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);
         }

server/src/tests/mod.rs

@@ -48,7 +48,7 @@ async fn start_test_server(port: u16, db: Option<CoreDB>) -> 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

tdb-macros/src/lib.rs

@@ -61,7 +61,8 @@ fn parse_dbtest(mut input: syn::ItemFn, rand: u16) -> Result<TokenStream, syn::E
     let body = quote! {
         let asyncdb = crate::coredb::CoreDB::new(
             crate::config::BGSave::Disabled,
-            crate::config::SnapshotConfig::default()
+            crate::config::SnapshotConfig::default(),
+            None
         ).unwrap();
         let addr = crate::tests::start_test_server(#rand, Some(asyncdb.clone())).await;
         let mut stream = tokio::net::TcpStream::connect(&addr).await.unwrap();