Move ArchiveFormat into own module (#23562)

runtime/src/snapshot_utils.rs

@@ -43,6 +43,9 @@ use {
     thiserror::Error,
 };
 
+mod archive_format;
+pub use archive_format::*;
+
 pub const SNAPSHOT_STATUS_CACHE_FILENAME: &str = "status_cache";
 pub const DEFAULT_FULL_SNAPSHOT_ARCHIVE_INTERVAL_SLOTS: Slot = 25_000;
 pub const DEFAULT_INCREMENTAL_SNAPSHOT_ARCHIVE_INTERVAL_SLOTS: Slot = 100;
@@ -112,15 +115,6 @@ impl SnapshotVersion {
     }
 }
 
-/// The different archive formats used for snapshots
-#[derive(Copy, Clone, Debug, Eq, PartialEq)]
-pub enum ArchiveFormat {
-    TarBzip2,
-    TarGzip,
-    TarZstd,
-    Tar,
-}
-
 /// A slot and the path to its bank snapshot
 #[derive(PartialEq, Eq, Debug)]
 pub struct BankSnapshotInfo {
@@ -209,15 +203,6 @@ pub enum SnapshotError {
 }
 pub type Result<T> = std::result::Result<T, SnapshotError>;
 
-fn get_archive_ext(archive_format: ArchiveFormat) -> &'static str {
-    match archive_format {
-        ArchiveFormat::TarBzip2 => "tar.bz2",
-        ArchiveFormat::TarGzip => "tar.gz",
-        ArchiveFormat::TarZstd => "tar.zst",
-        ArchiveFormat::Tar => "tar",
-    }
-}
-
 /// If the validator halts in the middle of `archive_snapshot_package()`, the temporary staging
 /// directory won't be cleaned up.  Call this function to clean them up.
 pub fn remove_tmp_snapshot_archives(snapshot_archives_dir: impl AsRef<Path>) {
@@ -322,14 +307,12 @@ pub fn archive_snapshot_package(
             .map_err(|e| SnapshotError::IoWithSource(e, "write version file"))?;
     }
 
-    let file_ext = get_archive_ext(snapshot_package.archive_format());
-
     // Tar the staging directory into the archive at `archive_path`
     let archive_path = tar_dir.join(format!(
         "{}{}.{}",
         staging_dir_prefix,
         snapshot_package.slot(),
-        file_ext
+        snapshot_package.archive_format().extension(),
     ));
 
     {
@@ -1062,7 +1045,7 @@ pub fn build_full_snapshot_archive_path(
         "snapshot-{}-{}.{}",
         slot,
         hash,
-        get_archive_ext(archive_format),
+        archive_format.extension(),
     ))
 }
 
@@ -1081,20 +1064,10 @@ pub fn build_incremental_snapshot_archive_path(
         base_slot,
         slot,
         hash,
-        get_archive_ext(archive_format),
+        archive_format.extension(),
     ))
 }
 
-fn archive_format_from_str(archive_format: &str) -> Option<ArchiveFormat> {
-    match archive_format {
-        "tar.bz2" => Some(ArchiveFormat::TarBzip2),
-        "tar.gz" => Some(ArchiveFormat::TarGzip),
-        "tar.zst" => Some(ArchiveFormat::TarZstd),
-        "tar" => Some(ArchiveFormat::Tar),
-        _ => None,
-    }
-}
-
 /// Parse a full snapshot archive filename into its Slot, Hash, and Archive Format
 pub(crate) fn parse_full_snapshot_archive_filename(
     archive_filename: &str,
@@ -1115,7 +1088,8 @@ pub(crate) fn parse_full_snapshot_archive_filename(
                 .ok()?;
             let archive_format = captures
                 .name("ext")
-                .map(|x| archive_format_from_str(x.as_str()))??;
+                .map(|x| x.as_str().parse::<ArchiveFormat>())?
+                .ok()?;
 
             Some((slot, hash, archive_format))
         })
@@ -1150,7 +1124,8 @@ pub(crate) fn parse_incremental_snapshot_archive_filename(
                 .ok()?;
             let archive_format = captures
                 .name("ext")
-                .map(|x| archive_format_from_str(x.as_str()))??;
+                .map(|x| x.as_str().parse::<ArchiveFormat>())?
+                .ok()?;
 
             Some((base_slot, slot, hash, archive_format))
         })

runtime/src/snapshot_utils/archive_format.rs

@@ -0,0 +1,118 @@
+use std::str::FromStr;
+
+pub const TAR_BZIP2_EXTENSION: &str = "tar.bz2";
+pub const TAR_GZIP_EXTENSION: &str = "tar.gz";
+pub const TAR_ZSTD_EXTENSION: &str = "tar.zst";
+pub const TAR_EXTENSION: &str = "tar";
+
+/// The different archive formats used for snapshots
+#[derive(Copy, Clone, Debug, Eq, PartialEq)]
+pub enum ArchiveFormat {
+    TarBzip2,
+    TarGzip,
+    TarZstd,
+    Tar,
+}
+
+impl ArchiveFormat {
+    /// Get the file extension for the ArchiveFormat
+    pub fn extension(&self) -> &str {
+        match self {
+            ArchiveFormat::TarBzip2 => TAR_BZIP2_EXTENSION,
+            ArchiveFormat::TarGzip => TAR_GZIP_EXTENSION,
+            ArchiveFormat::TarZstd => TAR_ZSTD_EXTENSION,
+            ArchiveFormat::Tar => TAR_EXTENSION,
+        }
+    }
+}
+
+// Change this to `impl<S: AsRef<str>> TryFrom<S> for ArchiveFormat [...]`
+// once this Rust bug is fixed: https://github.com/rust-lang/rust/issues/50133
+impl TryFrom<&str> for ArchiveFormat {
+    type Error = ParseError;
+
+    fn try_from(extension: &str) -> Result<Self, Self::Error> {
+        match extension {
+            TAR_BZIP2_EXTENSION => Ok(ArchiveFormat::TarBzip2),
+            TAR_GZIP_EXTENSION => Ok(ArchiveFormat::TarGzip),
+            TAR_ZSTD_EXTENSION => Ok(ArchiveFormat::TarZstd),
+            TAR_EXTENSION => Ok(ArchiveFormat::Tar),
+            _ => Err(ParseError::InvalidExtension),
+        }
+    }
+}
+
+impl FromStr for ArchiveFormat {
+    type Err = ParseError;
+
+    fn from_str(extension: &str) -> Result<Self, Self::Err> {
+        Self::try_from(extension)
+    }
+}
+
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub enum ParseError {
+    InvalidExtension,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    const INVALID_EXTENSION: &str = "zip";
+
+    #[test]
+    fn test_extension() {
+        assert_eq!(ArchiveFormat::TarBzip2.extension(), TAR_BZIP2_EXTENSION);
+        assert_eq!(ArchiveFormat::TarGzip.extension(), TAR_GZIP_EXTENSION);
+        assert_eq!(ArchiveFormat::TarZstd.extension(), TAR_ZSTD_EXTENSION);
+        assert_eq!(ArchiveFormat::Tar.extension(), TAR_EXTENSION);
+    }
+
+    #[test]
+    fn test_try_from() {
+        assert_eq!(
+            ArchiveFormat::try_from(TAR_BZIP2_EXTENSION),
+            Ok(ArchiveFormat::TarBzip2)
+        );
+        assert_eq!(
+            ArchiveFormat::try_from(TAR_GZIP_EXTENSION),
+            Ok(ArchiveFormat::TarGzip)
+        );
+        assert_eq!(
+            ArchiveFormat::try_from(TAR_ZSTD_EXTENSION),
+            Ok(ArchiveFormat::TarZstd)
+        );
+        assert_eq!(
+            ArchiveFormat::try_from(TAR_EXTENSION),
+            Ok(ArchiveFormat::Tar)
+        );
+        assert_eq!(
+            ArchiveFormat::try_from(INVALID_EXTENSION),
+            Err(ParseError::InvalidExtension)
+        );
+    }
+
+    #[test]
+    fn test_from_str() {
+        assert_eq!(
+            ArchiveFormat::from_str(TAR_BZIP2_EXTENSION),
+            Ok(ArchiveFormat::TarBzip2)
+        );
+        assert_eq!(
+            ArchiveFormat::from_str(TAR_GZIP_EXTENSION),
+            Ok(ArchiveFormat::TarGzip)
+        );
+        assert_eq!(
+            ArchiveFormat::from_str(TAR_ZSTD_EXTENSION),
+            Ok(ArchiveFormat::TarZstd)
+        );
+        assert_eq!(
+            ArchiveFormat::from_str(TAR_EXTENSION),
+            Ok(ArchiveFormat::Tar)
+        );
+        assert_eq!(
+            ArchiveFormat::from_str(INVALID_EXTENSION),
+            Err(ParseError::InvalidExtension)
+        );
+    }
+}