feat: add back the file size and count measurementst

Sponsored-by: author

1 files changed, 38 insertions, 2 deletions

diff --git a/src/suite.rs b/src/suite.rs
index 4e35723..2e4cede 100644
--- a/src/suite.rs
+++ b/src/suite.rs
@@ -3,9 +3,10 @@ use crate::specification::{Create, FileCount};
 use crate::step::Step;
 use log::{debug, info};
 use std::fs::File;
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::time::Instant;
 use tempfile::{tempdir, TempDir};
+use walkdir::WalkDir;
 
 /// A running benchmark suite.
 ///
@@ -25,6 +26,14 @@ pub enum SuiteError {
     /// File creation failed.
     #[error("Failed to create file {0}: {1}")]
     CreateFile(PathBuf, std::io::Error),
+
+    /// Error from counting files.
+    #[error("Error counting files in {0}: {1}")]
+    FileCount(PathBuf, walkdir::Error),
+
+    /// Error looking up file metadata.
+    #[error("Error looking up file metadata: {0}: {1}")]
+    FileMeta(PathBuf, walkdir::Error),
 }
 
 impl Suite {
@@ -106,6 +115,10 @@ impl Benchmark {
         self.tempdir.as_ref().unwrap()
     }
 
+    fn path(&self) -> &Path {
+        self.tempdir().path()
+    }
+
     fn start(&mut self) -> Result<OpMeasurements, SuiteError> {
         info!("starting benchmark {}", self.name());
         Ok(OpMeasurements::new(self.name(), Operation::Start))
@@ -143,7 +156,11 @@ impl Benchmark {
 
     fn backup(&mut self, n: usize) -> Result<OpMeasurements, SuiteError> {
         info!("making backup {} in benchmark {}", n, self.name());
-        Ok(OpMeasurements::new(self.name(), Operation::Backup))
+        let mut om = OpMeasurements::new(self.name(), Operation::Backup);
+        let stats = filestats(self.path())?;
+        om.push(Measurement::TotalFiles(stats.count));
+        om.push(Measurement::TotalData(stats.size));
+        Ok(om)
     }
 
     fn restore(&mut self, n: usize) -> Result<OpMeasurements, SuiteError> {
@@ -151,3 +168,22 @@ impl Benchmark {
         Ok(OpMeasurements::new(self.name(), Operation::Restore))
     }
 }
+
+#[derive(Debug, Default)]
+struct FileStats {
+    count: u64,
+    size: u64,
+}
+
+fn filestats(dirname: &Path) -> Result<FileStats, SuiteError> {
+    let mut stats = FileStats::default();
+    for e in WalkDir::new(dirname) {
+        let e = e.map_err(|err| SuiteError::FileCount(dirname.to_path_buf(), err))?;
+        stats.count += 1;
+        stats.size += e
+            .metadata()
+            .map_err(|err| SuiteError::FileMeta(e.path().to_path_buf(), err))?
+            .len();
+    }
+    Ok(stats)
+}