Adds zstd compression level to SnapshotConfig (solana-labs#4554)

brooksprumo · Jan 21, 2025 · 018699f · 018699f
1 parent 4fc75a1
commit 018699f
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 23 deletions.
diff --git a/download-utils/src/lib.rs b/download-utils/src/lib.rs
@@ -5,7 +5,7 @@ use {
     solana_runtime::{
         snapshot_hash::SnapshotHash,
         snapshot_package::SnapshotKind,
-        snapshot_utils::{self, ArchiveFormat},
+        snapshot_utils::{self, ArchiveFormat, ZstdConfig},
     },
     solana_sdk::{clock::Slot, genesis_config::DEFAULT_GENESIS_ARCHIVE},
     std::{
@@ -67,7 +67,9 @@ pub fn download_snapshot_archive(
     fs::create_dir_all(&snapshot_archives_remote_dir).unwrap();
 
     for archive_format in [
-        ArchiveFormat::TarZstd,
+        ArchiveFormat::TarZstd {
+            config: ZstdConfig::default(),
+        },
         ArchiveFormat::TarGzip,
         ArchiveFormat::TarBzip2,
         ArchiveFormat::TarLz4,

diff --git a/runtime/src/snapshot_bank_utils.rs b/runtime/src/snapshot_bank_utils.rs
@@ -1150,7 +1150,7 @@ mod tests {
                 SnapshotVersion::default(),
                 &snapshot_archives_dir,
                 &snapshot_archives_dir,
-                ArchiveFormat::TarZstd,
+                ArchiveFormat::Tar,
             )
             .unwrap();
 
@@ -1395,7 +1395,7 @@ mod tests {
         let bank_snapshots_dir = tempfile::TempDir::new().unwrap();
         let full_snapshot_archives_dir = tempfile::TempDir::new().unwrap();
         let incremental_snapshot_archives_dir = tempfile::TempDir::new().unwrap();
-        let snapshot_archive_format = ArchiveFormat::TarZstd;
+        let snapshot_archive_format = ArchiveFormat::Tar;
 
         let full_snapshot_slot = slot;
         let full_snapshot_archive_info = bank_to_full_snapshot_archive(
@@ -1878,7 +1878,7 @@ mod tests {
             SnapshotVersion::default(),
             &snapshot_archives_dir,
             &snapshot_archives_dir,
-            ArchiveFormat::TarZstd,
+            ArchiveFormat::Tar,
         )
         .unwrap();
 

diff --git a/runtime/src/snapshot_config.rs b/runtime/src/snapshot_config.rs
@@ -1,7 +1,7 @@
 use {
     crate::{
         snapshot_bank_utils,
-        snapshot_utils::{self, ArchiveFormat, SnapshotVersion},
+        snapshot_utils::{self, ArchiveFormat, SnapshotVersion, ZstdConfig},
     },
     solana_sdk::clock::Slot,
     std::{num::NonZeroUsize, path::PathBuf},
@@ -59,7 +59,9 @@ impl Default for SnapshotConfig {
             full_snapshot_archives_dir: PathBuf::default(),
             incremental_snapshot_archives_dir: PathBuf::default(),
             bank_snapshots_dir: PathBuf::default(),
-            archive_format: ArchiveFormat::TarZstd,
+            archive_format: ArchiveFormat::TarZstd {
+                config: ZstdConfig::default(),
+            },
             snapshot_version: SnapshotVersion::default(),
             maximum_full_snapshot_archives_to_retain:
                 snapshot_utils::DEFAULT_MAX_FULL_SNAPSHOT_ARCHIVES_TO_RETAIN,

diff --git a/runtime/src/snapshot_utils.rs b/runtime/src/snapshot_utils.rs
@@ -1086,10 +1086,10 @@ fn archive_snapshot(
                 do_archive_files(&mut encoder)?;
                 encoder.finish().map_err(E::FinishEncoder)?;
             }
-            ArchiveFormat::TarZstd => {
-                // Compression level of 1 is optimized for speed.
+            ArchiveFormat::TarZstd { config } => {
                 let mut encoder =
-                    zstd::stream::Encoder::new(archive_file, 1).map_err(E::CreateEncoder)?;
+                    zstd::stream::Encoder::new(archive_file, config.compression_level)
+                        .map_err(E::CreateEncoder)?;
                 do_archive_files(&mut encoder)?;
                 encoder.finish().map_err(E::FinishEncoder)?;
             }
@@ -2270,7 +2270,7 @@ fn untar_snapshot_create_shared_buffer(
     match archive_format {
         ArchiveFormat::TarBzip2 => SharedBuffer::new(BzDecoder::new(BufReader::new(open_file()))),
         ArchiveFormat::TarGzip => SharedBuffer::new(GzDecoder::new(BufReader::new(open_file()))),
-        ArchiveFormat::TarZstd => SharedBuffer::new(
+        ArchiveFormat::TarZstd { .. } => SharedBuffer::new(
             zstd::stream::read::Decoder::new(BufReader::new(open_file())).unwrap(),
         ),
         ArchiveFormat::TarLz4 => {
@@ -2738,7 +2738,13 @@ mod tests {
                 Hash::default()
             ))
             .unwrap(),
-            (43, SnapshotHash(Hash::default()), ArchiveFormat::TarZstd)
+            (
+                43,
+                SnapshotHash(Hash::default()),
+                ArchiveFormat::TarZstd {
+                    config: ZstdConfig::default(),
+                }
+            )
         );
         assert_eq!(
             parse_full_snapshot_archive_filename(&format!("snapshot-44-{}.tar", Hash::default()))
@@ -2819,7 +2825,9 @@ mod tests {
                 43,
                 234,
                 SnapshotHash(Hash::default()),
-                ArchiveFormat::TarZstd
+                ArchiveFormat::TarZstd {
+                    config: ZstdConfig::default(),
+                }
             )
         );
         assert_eq!(

diff --git a/runtime/src/snapshot_utils/archive_format.rs b/runtime/src/snapshot_utils/archive_format.rs
@@ -23,7 +23,7 @@ pub const TAR_EXTENSION: &str = "tar";
 pub enum ArchiveFormat {
     TarBzip2,
     TarGzip,
-    TarZstd,
+    TarZstd { config: ZstdConfig },
     TarLz4,
     Tar,
 }
@@ -34,15 +34,17 @@ impl ArchiveFormat {
         match self {
             ArchiveFormat::TarBzip2 => TAR_BZIP2_EXTENSION,
             ArchiveFormat::TarGzip => TAR_GZIP_EXTENSION,
-            ArchiveFormat::TarZstd => TAR_ZSTD_EXTENSION,
+            ArchiveFormat::TarZstd { .. } => TAR_ZSTD_EXTENSION,
             ArchiveFormat::TarLz4 => TAR_LZ4_EXTENSION,
             ArchiveFormat::Tar => TAR_EXTENSION,
         }
     }
 
     pub fn from_cli_arg(archive_format_str: &str) -> Option<ArchiveFormat> {
         match archive_format_str {
-            "zstd" => Some(ArchiveFormat::TarZstd),
+            "zstd" => Some(ArchiveFormat::TarZstd {
+                config: ZstdConfig::default(),
+            }),
             "lz4" => Some(ArchiveFormat::TarLz4),
             _ => None,
         }
@@ -58,7 +60,9 @@ impl TryFrom<&str> for ArchiveFormat {
         match extension {
             TAR_BZIP2_EXTENSION => Ok(ArchiveFormat::TarBzip2),
             TAR_GZIP_EXTENSION => Ok(ArchiveFormat::TarGzip),
-            TAR_ZSTD_EXTENSION => Ok(ArchiveFormat::TarZstd),
+            TAR_ZSTD_EXTENSION => Ok(ArchiveFormat::TarZstd {
+                config: ZstdConfig::default(),
+            }),
             TAR_LZ4_EXTENSION => Ok(ArchiveFormat::TarLz4),
             TAR_EXTENSION => Ok(ArchiveFormat::Tar),
             _ => Err(ParseError::InvalidExtension(extension.to_string())),
@@ -89,6 +93,13 @@ impl fmt::Display for ParseError {
     }
 }
 
+/// Configuration when using zstd as the snapshot archive format
+#[derive(Debug, Default, Copy, Clone, Eq, PartialEq)]
+pub struct ZstdConfig {
+    /// The compression level to use when archiving with zstd
+    pub compression_level: i32,
+}
+
 #[cfg(test)]
 mod tests {
     use {super::*, std::iter::zip};
@@ -98,7 +109,13 @@ mod tests {
     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::TarZstd {
+                config: ZstdConfig::default(),
+            }
+            .extension(),
+            TAR_ZSTD_EXTENSION
+        );
         assert_eq!(ArchiveFormat::TarLz4.extension(), TAR_LZ4_EXTENSION);
         assert_eq!(ArchiveFormat::Tar.extension(), TAR_EXTENSION);
     }
@@ -115,7 +132,9 @@ mod tests {
         );
         assert_eq!(
             ArchiveFormat::try_from(TAR_ZSTD_EXTENSION),
-            Ok(ArchiveFormat::TarZstd)
+            Ok(ArchiveFormat::TarZstd {
+                config: ZstdConfig::default(),
+            })
         );
         assert_eq!(
             ArchiveFormat::try_from(TAR_LZ4_EXTENSION),
@@ -143,7 +162,9 @@ mod tests {
         );
         assert_eq!(
             ArchiveFormat::from_str(TAR_ZSTD_EXTENSION),
-            Ok(ArchiveFormat::TarZstd)
+            Ok(ArchiveFormat::TarZstd {
+                config: ZstdConfig::default(),
+            })
         );
         assert_eq!(
             ArchiveFormat::from_str(TAR_LZ4_EXTENSION),
@@ -161,7 +182,12 @@ mod tests {
 
     #[test]
     fn test_from_cli_arg() {
-        let golden = [Some(ArchiveFormat::TarZstd), Some(ArchiveFormat::TarLz4)];
+        let golden = [
+            Some(ArchiveFormat::TarZstd {
+                config: ZstdConfig::default(),
+            }),
+            Some(ArchiveFormat::TarLz4),
+        ];
 
         for (arg, expected) in zip(SUPPORTED_ARCHIVE_COMPRESSION.iter(), golden.into_iter()) {
             assert_eq!(ArchiveFormat::from_cli_arg(arg), expected);

diff --git a/validator/src/main.rs b/validator/src/main.rs
@@ -1730,8 +1730,13 @@ pub fn main() {
 
     let archive_format = {
         let archive_format_str = value_t_or_exit!(matches, "snapshot_archive_format", String);
-        ArchiveFormat::from_cli_arg(&archive_format_str)
-            .unwrap_or_else(|| panic!("Archive format not recognized: {archive_format_str}"))
+        let mut archive_format = ArchiveFormat::from_cli_arg(&archive_format_str)
+            .unwrap_or_else(|| panic!("Archive format not recognized: {archive_format_str}"));
+        if let ArchiveFormat::TarZstd { config } = &mut archive_format {
+            // level 1 is optimized for speed
+            config.compression_level = 1;
+        }
+        archive_format
     };
 
     let snapshot_version =