Clean up technical debt

src/common.rs

@@ -1 +1,3 @@
 pub mod db;
+pub mod lmdb_utils;
+pub mod progress;

src/common/db.rs

@@ -39,6 +39,7 @@ use thiserror::Error;
 
 use casper_types::bytesrepr::Error as BytesreprError;
 
+pub const STORAGE_FILE_NAME: &str = "storage.lmdb";
 const ENTRY_LOG_INTERVAL: usize = 100_000;
 const MAX_DB_READERS: u32 = 100;
 

src/common/db/tests.rs

@@ -17,9 +17,9 @@ fn gen_faulty_bytes(rng: &mut ThreadRng) -> Vec<u8> {
 
 fn populate_db(env: &Environment, db: &LmdbDatabase) {
     let mut rng = rand::thread_rng();
-    let entries_count = rng.gen_range(10u32..100u32);
+    let entry_count = rng.gen_range(10u32..100u32);
     let mut rw_tx = env.begin_rw_txn().expect("couldn't begin rw transaction");
-    for i in 0..entries_count {
+    for i in 0..entry_count {
         let bytes = gen_bytes(&mut rng);
         let key: [u8; 4] = i.to_le_bytes();
         rw_tx.put(*db, &key, &bytes, WriteFlags::empty()).unwrap();
@@ -29,9 +29,9 @@ fn populate_db(env: &Environment, db: &LmdbDatabase) {
 
 fn populate_faulty_db(env: &Environment, db: &LmdbDatabase) {
     let mut rng = rand::thread_rng();
-    let entries_count = rng.gen_range(10u32..100u32);
+    let entry_count = rng.gen_range(10u32..100u32);
     let mut rw_tx = env.begin_rw_txn().expect("couldn't begin rw transaction");
-    for i in 0..entries_count {
+    for i in 0..entry_count {
         let bytes = if i % 5 == 0 {
             gen_faulty_bytes(&mut rng)
         } else {
@@ -141,7 +141,7 @@ fn sanity_check_ser_deser() {
 
 #[test]
 fn good_db_should_pass_check() {
-    let fixture = LmdbTestFixture::new(Some(MockDb::db_name()));
+    let fixture = LmdbTestFixture::new(Some(MockDb::db_name()), None);
     populate_db(&fixture.env, &fixture.db);
 
     assert!(MockDb::check_db(&fixture.env, true, 0).is_ok());
@@ -152,7 +152,7 @@ fn good_db_should_pass_check() {
 
 #[test]
 fn bad_db_should_fail_check() {
-    let fixture = LmdbTestFixture::new(Some(MockDb::db_name()));
+    let fixture = LmdbTestFixture::new(Some(MockDb::db_name()), None);
     populate_faulty_db(&fixture.env, &fixture.db);
 
     assert!(MockDb::check_db(&fixture.env, true, 0).is_err());

src/common/lmdb_utils.rs

@@ -0,0 +1,90 @@
+use std::result::Result;
+
+use lmdb::{Database, Error, Transaction};
+use lmdb_sys::{mdb_stat, MDB_stat};
+
+/// Retrieves the number of entries in a database.
+pub fn entry_count<T: Transaction>(txn: &'_ T, database: Database) -> Result<usize, Error> {
+    let mut stat = MDB_stat {
+        ms_psize: 0,
+        ms_depth: 0,
+        ms_branch_pages: 0,
+        ms_leaf_pages: 0,
+        ms_overflow_pages: 0,
+        ms_entries: 0,
+    };
+    let result = unsafe { mdb_stat(txn.txn(), database.dbi(), &mut stat as *mut MDB_stat) };
+    if result != 0 {
+        Err(Error::from_err_code(result))
+    } else {
+        Ok(stat.ms_entries)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use lmdb::{Transaction, WriteFlags};
+
+    use crate::test_utils::LmdbTestFixture;
+
+    use super::entry_count;
+
+    #[test]
+    fn db_entry_count() {
+        let fixture = LmdbTestFixture::new(None, None);
+        let env = &fixture.env;
+        let db = fixture.db;
+
+        if let Ok(txn) = env.begin_ro_txn() {
+            assert_eq!(entry_count(&txn, db).unwrap(), 0);
+            txn.commit().unwrap();
+        }
+
+        let first_dummy_input = [0u8, 1u8];
+        let second_dummy_input = [1u8, 2u8];
+        // Insert the first entry into the database.
+        if let Ok(mut txn) = env.begin_rw_txn() {
+            txn.put(
+                fixture.db,
+                &first_dummy_input,
+                &first_dummy_input,
+                WriteFlags::empty(),
+            )
+            .unwrap();
+            txn.commit().unwrap();
+        };
+
+        if let Ok(txn) = env.begin_ro_txn() {
+            assert_eq!(entry_count(&txn, db).unwrap(), 1);
+            txn.commit().unwrap();
+        }
+
+        // Insert the second entry into the database.
+        if let Ok(mut txn) = env.begin_rw_txn() {
+            txn.put(
+                fixture.db,
+                &second_dummy_input,
+                &second_dummy_input,
+                WriteFlags::empty(),
+            )
+            .unwrap();
+            txn.commit().unwrap();
+        };
+
+        if let Ok(txn) = env.begin_ro_txn() {
+            assert_eq!(entry_count(&txn, db).unwrap(), 2);
+            txn.commit().unwrap();
+        };
+
+        // Delete the first entry from the database.
+        if let Ok(mut txn) = env.begin_rw_txn() {
+            txn.del(fixture.db, &first_dummy_input, None).unwrap();
+            txn.commit().unwrap();
+        };
+
+        if let Ok(txn) = env.begin_ro_txn() {
+            assert_eq!(entry_count(&txn, db).unwrap(), 1);
+            txn.commit().unwrap();
+        };
+    }
+}

src/common/progress.rs

@@ -0,0 +1,71 @@
+use std::result::Result;
+
+use log::warn;
+
+// The tracker will log 20 times during its operation.
+const STEPS: usize = 20;
+// The tracker will log progress every in 5% completion intervals.
+const PROGRESS_MULTIPLIER: u64 = 100 / STEPS as u64;
+// Error message for initialization of a progress tracker with nothing
+// to process.
+const NULL_TOTAL_TO_PROCESS_ERROR: &str = "Cannot initialize total to process with 0";
+
+/// Tracks and logs progress of an operation in a human readable form.
+/// Whenever (1 / number of steps) of the total amount has been processed,
+/// this structure calls the `log_progress` function, which takes the
+/// percentage completed so far as a parameter.
+pub struct ProgressTracker {
+    /// Total amount there is to process.
+    total_to_process: usize,
+    /// Amount processed so far.
+    processed: usize,
+    /// Internal counter to keep track of the number of steps completed
+    /// so far relative to the maximum amount of steps this operation
+    /// will do, defined in `STEPS`.
+    progress_factor: u64,
+    /// Function which takes the completion rate as a percentage as
+    /// input. It is called zero or more times as progress is being made
+    /// using `ProgressTracker::advance_by`. The purpose of this function
+    /// is to allow users to create custom log messages for their specific
+    /// operation.
+    log_progress: Box<dyn Fn(u64)>,
+}
+
+impl ProgressTracker {
+    /// Create a new progress tracker by initializing it with a non-zero
+    /// amount to be processed and a log function.
+    pub fn new(
+        total_to_process: usize,
+        log_progress: Box<dyn Fn(u64)>,
+    ) -> Result<Self, &'static str> {
+        if total_to_process == 0 {
+            Err(NULL_TOTAL_TO_PROCESS_ERROR)
+        } else {
+            Ok(Self {
+                total_to_process,
+                processed: 0,
+                progress_factor: 1,
+                log_progress,
+            })
+        }
+    }
+
+    /// Advance the progress tracker by a specific amount. If it passes
+    /// a milestone ((1 / STEP) of the total amount to process),
+    /// `log_progress` will be called with the current completion rate
+    /// as input.
+    pub fn advance_by(&mut self, step: usize) {
+        self.processed += step;
+        while self.processed >= (self.total_to_process * self.progress_factor as usize) / STEPS {
+            (*self.log_progress)(self.progress_factor * PROGRESS_MULTIPLIER);
+            self.progress_factor += 1;
+        }
+        if self.processed > self.total_to_process {
+            warn!(
+                "Exceeded total amount to process {} by {}",
+                self.total_to_process,
+                self.processed - self.total_to_process
+            );
+        }
+    }
+}

src/main.rs

@@ -9,7 +9,7 @@ use std::{fs::OpenOptions, process};
 use clap::{crate_description, crate_version, Arg, Command};
 use log::error;
 
-use subcommands::{archive, check, latest_block_summary, trie_compact, unsparse};
+use subcommands::{archive, check, latest_block_summary, trie_compact, unsparse, Error};
 
 const LOGGING: &str = "logging";
 
@@ -61,20 +61,26 @@ fn main() {
     );
 
     let (subcommand_name, matches) = arg_matches.subcommand().unwrap_or_else(|| {
-        error!("{}", cli().get_long_about().unwrap());
+        error!(
+            "{}",
+            cli().get_long_about().expect("should have long about")
+        );
         process::exit(1);
     });
 
-    let succeeded = match subcommand_name {
-        archive::COMMAND_NAME => archive::run(matches),
-        check::COMMAND_NAME => check::run(matches),
-        latest_block_summary::COMMAND_NAME => latest_block_summary::run(matches),
-        trie_compact::COMMAND_NAME => trie_compact::run(matches),
-        unsparse::COMMAND_NAME => unsparse::run(matches),
+    let result: Result<(), Error> = match subcommand_name {
+        archive::COMMAND_NAME => archive::run(matches).map_err(Error::from),
+        check::COMMAND_NAME => check::run(matches).map_err(Error::from),
+        latest_block_summary::COMMAND_NAME => {
+            latest_block_summary::run(matches).map_err(Error::from)
+        }
+        trie_compact::COMMAND_NAME => trie_compact::run(matches).map_err(Error::from),
+        unsparse::COMMAND_NAME => unsparse::run(matches).map_err(Error::from),
         _ => unreachable!("{} should be handled above", subcommand_name),
     };
 
-    if !succeeded {
+    if let Err(run_err) = result {
+        error!("{}", run_err);
         process::exit(1);
     }
 }

src/subcommands.rs

@@ -3,3 +3,27 @@ pub mod check;
 pub mod latest_block_summary;
 pub mod trie_compact;
 pub mod unsparse;
+
+use thiserror::Error as ThisError;
+
+use archive::{CreateError, UnpackError};
+use check::Error as CheckError;
+use latest_block_summary::Error as LatestBlockSummaryError;
+use trie_compact::Error as TrieCompactError;
+use unsparse::Error as UnsparseError;
+
+#[derive(ThisError, Debug)]
+pub enum Error {
+    #[error("Archive create failed: {0}")]
+    ArchiveCreate(#[from] CreateError),
+    #[error("Archive unpack failed: {0}")]
+    ArchiveUnpack(#[from] UnpackError),
+    #[error("Check command failed: {0}")]
+    Check(#[from] CheckError),
+    #[error("Latest block summary command failed: {0}")]
+    LatestBlockSummary(#[from] LatestBlockSummaryError),
+    #[error("Trie compact failed: {0}")]
+    TrieCompact(#[from] TrieCompactError),
+    #[error("Unsparse failed: {0}")]
+    Unsparse(#[from] UnsparseError),
+}

src/subcommands/archive.rs

@@ -1,6 +1,12 @@
 use std::process;
 
 use clap::{ArgMatches, Command};
+use thiserror::Error as ThisError;
+
+pub use create::Error as CreateError;
+pub use unpack::Error as UnpackError;
+
+use super::Error as SubcommandError;
 
 mod create;
 mod ring_buffer;
@@ -15,6 +21,23 @@ enum DisplayOrder {
     Unpack,
 }
 
+#[derive(ThisError, Debug)]
+pub enum Error {
+    #[error("create: {0}")]
+    Create(#[from] CreateError),
+    #[error("unpack: {0}")]
+    Unpack(#[from] UnpackError),
+}
+
+impl From<Error> for SubcommandError {
+    fn from(err: Error) -> Self {
+        match err {
+            Error::Create(create_err) => SubcommandError::ArchiveCreate(create_err),
+            Error::Unpack(unpack_err) => SubcommandError::ArchiveUnpack(unpack_err),
+        }
+    }
+}
+
 pub fn command(display_order: usize) -> Command<'static> {
     Command::new(COMMAND_NAME)
         .display_order(display_order)
@@ -23,14 +46,14 @@ pub fn command(display_order: usize) -> Command<'static> {
         .subcommand(unpack::command(DisplayOrder::Unpack as usize))
 }
 
-pub fn run(matches: &ArgMatches) -> bool {
+pub fn run(matches: &ArgMatches) -> Result<(), Error> {
     let (subcommand_name, matches) = matches.subcommand().unwrap_or_else(|| {
         process::exit(1);
     });
 
     match subcommand_name {
-        create::COMMAND_NAME => create::run(matches),
-        unpack::COMMAND_NAME => unpack::run(matches),
+        create::COMMAND_NAME => create::run(matches).map_err(Error::Create),
+        unpack::COMMAND_NAME => unpack::run(matches).map_err(Error::Unpack),
         _ => unreachable!("{} should be handled above", subcommand_name),
     }
 }

src/subcommands/archive/create.rs

@@ -11,6 +11,7 @@ use thiserror::Error as ThisError;
 use super::zstd_utils::Error as ZstdError;
 
 pub const COMMAND_NAME: &str = "create";
+const OVERWRITE: &str = "overwrite";
 const OUTPUT: &str = "output";
 const DB: &str = "db-dir";
 
@@ -29,6 +30,7 @@ pub enum Error {
 enum DisplayOrder {
     Db,
     Output,
+    Overwrite,
 }
 
 pub fn command(display_order: usize) -> Command<'static> {
@@ -57,16 +59,23 @@ pub fn command(display_order: usize) -> Command<'static> {
                 .value_name("FILE_PATH")
                 .help("Output file path for the compressed tar archive."),
         )
+        .arg(
+            Arg::new(OVERWRITE)
+                .display_order(DisplayOrder::Overwrite as usize)
+                .required(false)
+                .short('w')
+                .long(OVERWRITE)
+                .takes_value(false)
+                .help(
+                    "Overwrite an already existing archive file in destination \
+                    directory.",
+                ),
+        )
 }
 
-pub fn run(matches: &ArgMatches) -> bool {
+pub fn run(matches: &ArgMatches) -> Result<(), Error> {
     let db_path = matches.value_of(DB).unwrap();
     let dest = matches.value_of(OUTPUT).unwrap();
-    let result = pack::create_archive(db_path, dest);
-
-    if let Err(error) = &result {
-        error!("Archive packing failed. {}", error);
-    }
-
-    result.is_ok()
+    let overwrite = matches.is_present(OVERWRITE);
+    pack::create_archive(db_path, dest, overwrite)
 }