Merge pull request #1011 from dlachaumepalo/damien/983/clean-pending-snapshot-directory

Clean pending_snapshot directory of aggregator

Author: jpraynaud

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-aggregator"
-version = "0.3.40"
+version = "0.3.41"
 description = "A Mithril Aggregator server"
 authors = { workspace = true }
 edition = { workspace = true }

mithril-aggregator/Cargo.toml

@@ -81,7 +81,7 @@ impl CardanoImmutableFilesFullArtifactBuilder {
         let location = self
             .snapshot_uploader
             .upload_snapshot(ongoing_snapshot.get_file_path())
-            .await?;
+            .await;
 
         if let Err(error) = tokio::fs::remove_file(ongoing_snapshot.get_file_path()).await {
             warn!(
@@ -90,7 +90,7 @@ impl CardanoImmutableFilesFullArtifactBuilder {
             );
         }
 
-        Ok(vec![location])
+        Ok(vec![location?])
     }
 
     async fn create_snapshot(
@@ -152,7 +152,7 @@ mod tests {
 
     use super::*;
 
-    use crate::{DumbSnapshotUploader, DumbSnapshotter};
+    use crate::{snapshot_uploaders::MockSnapshotUploader, DumbSnapshotUploader, DumbSnapshotter};
 
     #[tokio::test]
     async fn should_compute_valid_artifact() {
@@ -249,4 +249,32 @@ mod tests {
             ongoing_snapshot.get_file_path()
         );
     }
+
+    #[tokio::test]
+    async fn remove_snapshot_archive_after_upload_even_if_an_error_occured() {
+        let file = NamedTempFile::new().unwrap();
+        let file_path = file.path();
+        let snapshot = OngoingSnapshot::new(file_path.to_path_buf(), 7331);
+        let mut snapshot_uploader = MockSnapshotUploader::new();
+        snapshot_uploader
+            .expect_upload_snapshot()
+            .return_once(|_| Err("an error".to_string()))
+            .once();
+
+        let cardano_immutable_files_full_artifact_builder =
+            CardanoImmutableFilesFullArtifactBuilder::new(
+                Arc::new(DumbSnapshotter::new()),
+                Arc::new(snapshot_uploader),
+            );
+
+        cardano_immutable_files_full_artifact_builder
+            .upload_snapshot_archive(&snapshot)
+            .await
+            .expect_err("Snapshot upload should have failed");
+
+        assert!(
+            !file_path.exists(),
+            "Ongoing snapshot file should have been removed even after upload failure"
+        );
+    }
 }

mithril-aggregator/src/artifact_builder/cardano_immutable_files_full.rs

@@ -574,23 +574,10 @@ impl DependenciesBuilder {
                     .snapshot_directory
                     .join("pending_snapshot");
 
-                // **TODO** this code should be in the snapshotter constructor.
-                if !ongoing_snapshot_directory.exists() {
-                    std::fs::create_dir(&ongoing_snapshot_directory).map_err(|e| {
-                        DependenciesBuilderError::Initialization {
-                            message: format!(
-                                "Cannot create snapshotter directory '{}'.",
-                                ongoing_snapshot_directory.display()
-                            ),
-                            error: Some(e.into()),
-                        }
-                    })?;
-                }
-
                 Arc::new(GzipSnapshotter::new(
                     self.configuration.db_directory.clone(),
                     ongoing_snapshot_directory,
-                ))
+                )?)
             }
             _ => Arc::new(DumbSnapshotter::new()),
         };

mithril-aggregator/src/dependency_injection/builder.rs

@@ -1,13 +1,16 @@
 use flate2::write::GzEncoder;
 use flate2::Compression;
-use slog_scope::info;
+use mithril_common::StdResult;
+use slog_scope::{info, warn};
 use std::error::Error as StdError;
 use std::fs::File;
 use std::io;
 use std::path::{Path, PathBuf};
 use std::sync::RwLock;
 use thiserror::Error;
 
+use crate::dependency_injection::DependenciesBuilderError;
+
 /// Define the ability to create snapshots.
 pub trait Snapshotter: Sync + Send {
     /// Create a new snapshot with the given archive name.
@@ -33,6 +36,7 @@ impl OngoingSnapshot {
     pub fn new(filepath: PathBuf, filesize: u64) -> Self {
         Self { filepath, filesize }
     }
+
     pub fn get_file_path(&self) -> &PathBuf {
         &self.filepath
     }
@@ -60,33 +64,69 @@ pub enum SnapshotError {
 
 impl Snapshotter for GzipSnapshotter {
     fn snapshot(&self, archive_name: &str) -> Result<OngoingSnapshot, SnapshotError> {
-        let filepath = self.create_archive(archive_name)?;
-        let filesize = std::fs::metadata(&filepath)
-            .map_err(|e| SnapshotError::GeneralError(e.to_string()))?
-            .len();
+        let archive_path = self.ongoing_snapshot_directory.join(archive_name);
+        let filesize = self.create_archive(&archive_path).map_err(|err| {
+            if archive_path.exists() {
+                if let Err(remove_error) = std::fs::remove_file(&archive_path) {
+                    warn!(
+                        " > Post snapshotter.snapshot failure, could not remove temporary archive at path: path:{}, err: {}",
+                        archive_path.display(),
+                        remove_error
+                    );
+                }
+            }
+
+            err
+        })?;
 
-        Ok(OngoingSnapshot { filepath, filesize })
+        Ok(OngoingSnapshot {
+            filepath: archive_path,
+            filesize,
+        })
     }
 }
 
 impl GzipSnapshotter {
     /// Snapshotter factory
-    pub fn new(db_directory: PathBuf, ongoing_snapshot_directory: PathBuf) -> Self {
-        Self {
+    pub fn new(
+        db_directory: PathBuf,
+        ongoing_snapshot_directory: PathBuf,
+    ) -> StdResult<GzipSnapshotter> {
+        if ongoing_snapshot_directory.exists() {
+            std::fs::remove_dir_all(&ongoing_snapshot_directory)?;
+        }
+
+        std::fs::create_dir(&ongoing_snapshot_directory).map_err(|e| {
+            DependenciesBuilderError::Initialization {
+                message: format!(
+                    "Cannot create snapshotter directory '{}'.",
+                    ongoing_snapshot_directory.display()
+                ),
+                error: Some(e.into()),
+            }
+        })?;
+
+        Ok(Self {
             db_directory,
             ongoing_snapshot_directory,
-        }
+        })
+    }
+
+    fn get_file_size(filepath: &Path) -> Result<u64, SnapshotError> {
+        let res = std::fs::metadata(filepath)
+            .map_err(|e| SnapshotError::GeneralError(e.to_string()))?
+            .len();
+        Ok(res)
     }
 
-    fn create_archive(&self, archive_name: &str) -> Result<PathBuf, SnapshotError> {
-        let path = self.ongoing_snapshot_directory.join(archive_name);
+    fn create_archive(&self, archive_path: &Path) -> Result<u64, SnapshotError> {
         info!(
             "compressing {} into {}",
             self.db_directory.display(),
-            path.display()
+            archive_path.display()
         );
 
-        let tar_gz = File::create(&path).map_err(SnapshotError::CreateArchiveError)?;
+        let tar_gz = File::create(archive_path).map_err(SnapshotError::CreateArchiveError)?;
         let enc = GzEncoder::new(tar_gz, Compression::default());
         let mut tar = tar::Builder::new(enc);
 
@@ -98,7 +138,9 @@ impl GzipSnapshotter {
             .map_err(SnapshotError::CreateArchiveError)?;
         gz.try_finish().map_err(SnapshotError::CreateArchiveError)?;
 
-        Ok(path)
+        let filesize = Self::get_file_size(archive_path)?;
+
+        Ok(filesize)
     }
 }
 
@@ -154,8 +196,23 @@ impl Snapshotter for DumbSnapshotter {
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
+
     use super::*;
 
+    fn get_test_directory(dir_name: &str) -> PathBuf {
+        let test_dir = std::env::temp_dir()
+            .join("mithril_test")
+            .join("snapshotter")
+            .join(dir_name);
+        if test_dir.exists() {
+            std::fs::remove_dir_all(&test_dir).unwrap();
+        }
+        std::fs::create_dir_all(&test_dir).unwrap();
+
+        test_dir
+    }
+
     #[test]
     fn test_dumb_snapshotter() {
         let snapshotter = DumbSnapshotter::new();
@@ -174,4 +231,62 @@ mod tests {
             )
         );
     }
+
+    #[test]
+    fn should_create_directory_if_does_not_exist() {
+        let test_dir = get_test_directory("should_create_directory_if_does_not_exist");
+        let pending_snapshot_directory = test_dir.join("pending_snapshot");
+        let db_directory = std::env::temp_dir().join("whatever");
+
+        Arc::new(GzipSnapshotter::new(db_directory, pending_snapshot_directory.clone()).unwrap());
+
+        assert!(pending_snapshot_directory.is_dir());
+    }
+
+    #[test]
+    fn should_clean_pending_snapshot_directory_if_already_exists() {
+        let test_dir =
+            get_test_directory("should_clean_pending_snapshot_directory_if_already_exists");
+        let pending_snapshot_directory = test_dir.join("pending_snapshot");
+        let db_directory = std::env::temp_dir().join("whatever");
+
+        std::fs::create_dir_all(&pending_snapshot_directory).unwrap();
+
+        std::fs::File::create(pending_snapshot_directory.join("whatever.txt")).unwrap();
+
+        Arc::new(GzipSnapshotter::new(db_directory, pending_snapshot_directory.clone()).unwrap());
+
+        assert_eq!(
+            0,
+            std::fs::read_dir(pending_snapshot_directory)
+                .unwrap()
+                .count()
+        );
+    }
+
+    #[test]
+    fn should_delete_tmp_file_in_pending_snapshot_directory_if_snapshotting_fail() {
+        let test_dir = get_test_directory(
+            "should_delete_tmp_file_in_pending_snapshot_directory_if_snapshotting_fail",
+        );
+        let pending_snapshot_directory = test_dir.join("pending_snapshot");
+        let db_directory = test_dir.join("db");
+
+        let snapshotter = Arc::new(
+            GzipSnapshotter::new(db_directory, pending_snapshot_directory.clone()).unwrap(),
+        );
+
+        // this file should not be deleted by the archive creation
+        std::fs::File::create(pending_snapshot_directory.join("other-process.file")).unwrap();
+
+        let _ = snapshotter
+            .snapshot("whatever.tar.gz")
+            .expect_err("Snapshotter::snapshot should fail if the db is empty.");
+        let remaining_files: Vec<String> = std::fs::read_dir(&pending_snapshot_directory)
+            .unwrap()
+            .map(|f| f.unwrap().file_name().to_str().unwrap().to_owned())
+            .collect();
+
+        assert_eq!(vec!["other-process.file".to_string()], remaining_files);
+    }
 }