From 1eddb6d1e945964adc615f0016020d0d30a0ed89 Mon Sep 17 00:00:00 2001 From: Brooks Prumo Date: Wed, 9 Mar 2022 16:09:34 -0600 Subject: [PATCH] Move ArchiveFormat into own module (#23562) --- runtime/src/snapshot_utils.rs | 45 ++----- runtime/src/snapshot_utils/archive_format.rs | 118 +++++++++++++++++++ 2 files changed, 128 insertions(+), 35 deletions(-) create mode 100644 runtime/src/snapshot_utils/archive_format.rs diff --git a/runtime/src/snapshot_utils.rs b/runtime/src/snapshot_utils.rs index 68b4efe22c..56d1889aa1 100644 --- a/runtime/src/snapshot_utils.rs +++ b/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 = std::result::Result; -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) { @@ -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 { - 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::())? + .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::())? + .ok()?; Some((base_slot, slot, hash, archive_format)) }) diff --git a/runtime/src/snapshot_utils/archive_format.rs b/runtime/src/snapshot_utils/archive_format.rs new file mode 100644 index 0000000000..ef0102e478 --- /dev/null +++ b/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> TryFrom 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 { + 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::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) + ); + } +}