Merge pull request #105 from rust-lang/dynamic-compression

pietroalbini · web-flow · commit 5254dbfd25d5 · 2020-12-10T17:54:47.000+01:00
Allow specifying the wanted compression formats at runtime
diff --git a/src/combiner.rs b/src/combiner.rs
@@ -1,8 +1,10 @@
 use super::Scripter;
 use super::Tarballer;
-use crate::util::*;
+use crate::{
+    compression::{CompressionFormat, CompressionFormats},
+    util::*,
+};
 use anyhow::{bail, Context, Result};
-use flate2::read::GzDecoder;
 use std::io::{Read, Write};
 use std::path::Path;
 use tar::Archive;
@@ -36,6 +38,9 @@ actor! {
 
         /// The location to put the final image and tarball.
         output_dir: String = "./dist",
+
+        /// The formats used to compress the tarball
+        compression_formats: CompressionFormats = CompressionFormats::default(),
     }
 }
 
@@ -59,15 +64,21 @@ impl Combiner {
             .filter(|s| !s.is_empty())
         {
             // Extract the input tarballs
-            let tar = GzDecoder::new(open_file(&input_tarball)?);
-            Archive::new(tar).unpack(&self.work_dir).with_context(|| {
-                format!(
-                    "unable to extract '{}' into '{}'",
-                    &input_tarball, self.work_dir
-                )
-            })?;
-
-            let pkg_name = input_tarball.trim_end_matches(".tar.gz");
+            let compression =
+                CompressionFormat::detect_from_path(input_tarball).ok_or_else(|| {
+                    anyhow::anyhow!("couldn't figure out the format of {}", input_tarball)
+                })?;
+            Archive::new(compression.decode(input_tarball)?)
+                .unpack(&self.work_dir)
+                .with_context(|| {
+                    format!(
+                        "unable to extract '{}' into '{}'",
+                        &input_tarball, self.work_dir
+                    )
+                })?;
+
+            let pkg_name =
+                input_tarball.trim_end_matches(&format!(".tar.{}", compression.extension()));
             let pkg_name = Path::new(pkg_name).file_name().unwrap();
             let pkg_dir = Path::new(&self.work_dir).join(&pkg_name);
 
@@ -121,7 +132,7 @@ impl Combiner {
             .rel_manifest_dir(self.rel_manifest_dir)
             .success_message(self.success_message)
             .legacy_manifest_dirs(self.legacy_manifest_dirs)
-            .output_script(path_to_str(&output_script)?);
+            .output_script(path_to_str(&output_script)?.into());
         scripter.run()?;
 
         // Make the tarballs.
@@ -131,7 +142,8 @@ impl Combiner {
         tarballer
             .work_dir(self.work_dir)
             .input(self.package_name)
-            .output(path_to_str(&output)?);
+            .output(path_to_str(&output)?.into())
+            .compression_formats(self.compression_formats.clone());
         tarballer.run()?;
 
         Ok(())
diff --git a/src/compression.rs b/src/compression.rs
@@ -0,0 +1,154 @@
+use anyhow::{Context, Error};
+use flate2::{read::GzDecoder, write::GzEncoder};
+use rayon::prelude::*;
+use std::{convert::TryFrom, io::Read, io::Write, path::Path};
+use xz2::{read::XzDecoder, write::XzEncoder};
+
+#[derive(Debug, Copy, Clone)]
+pub enum CompressionFormat {
+    Gz,
+    Xz,
+}
+
+impl CompressionFormat {
+    pub(crate) fn detect_from_path(path: impl AsRef<Path>) -> Option<Self> {
+        match path.as_ref().extension().and_then(|e| e.to_str()) {
+            Some("gz") => Some(CompressionFormat::Gz),
+            Some("xz") => Some(CompressionFormat::Xz),
+            _ => None,
+        }
+    }
+
+    pub(crate) fn extension(&self) -> &'static str {
+        match self {
+            CompressionFormat::Gz => "gz",
+            CompressionFormat::Xz => "xz",
+        }
+    }
+
+    pub(crate) fn encode(&self, path: impl AsRef<Path>) -> Result<Box<dyn Encoder>, Error> {
+        let mut os = path.as_ref().as_os_str().to_os_string();
+        os.push(format!(".{}", self.extension()));
+        let path = Path::new(&os);
+
+        if path.exists() {
+            crate::util::remove_file(path)?;
+        }
+        let file = crate::util::create_new_file(path)?;
+
+        Ok(match self {
+            CompressionFormat::Gz => Box::new(GzEncoder::new(file, flate2::Compression::best())),
+            CompressionFormat::Xz => {
+                // Note that preset 6 takes about 173MB of memory per thread, so we limit the number of
+                // threads to not blow out 32-bit hosts.  (We could be more precise with
+                // `MtStreamBuilder::memusage()` if desired.)
+                let stream = xz2::stream::MtStreamBuilder::new()
+                    .threads(Ord::min(num_cpus::get(), 8) as u32)
+                    .preset(6)
+                    .encoder()?;
+                Box::new(XzEncoder::new_stream(file, stream))
+            }
+        })
+    }
+
+    pub(crate) fn decode(&self, path: impl AsRef<Path>) -> Result<Box<dyn Read>, Error> {
+        let file = crate::util::open_file(path.as_ref())?;
+        Ok(match self {
+            CompressionFormat::Gz => Box::new(GzDecoder::new(file)),
+            CompressionFormat::Xz => Box::new(XzDecoder::new(file)),
+        })
+    }
+}
+
+/// This struct wraps Vec<CompressionFormat> in order to parse the value from the command line.
+#[derive(Debug, Clone)]
+pub struct CompressionFormats(Vec<CompressionFormat>);
+
+impl TryFrom<&'_ str> for CompressionFormats {
+    type Error = Error;
+
+    fn try_from(value: &str) -> Result<Self, Self::Error> {
+        let mut parsed = Vec::new();
+        for format in value.split(',') {
+            match format.trim() {
+                "gz" => parsed.push(CompressionFormat::Gz),
+                "xz" => parsed.push(CompressionFormat::Xz),
+                other => anyhow::bail!("unknown compression format: {}", other),
+            }
+        }
+        Ok(CompressionFormats(parsed))
+    }
+}
+
+impl Default for CompressionFormats {
+    fn default() -> Self {
+        Self(vec![CompressionFormat::Gz, CompressionFormat::Xz])
+    }
+}
+
+impl CompressionFormats {
+    pub(crate) fn iter(&self) -> impl Iterator<Item = CompressionFormat> + '_ {
+        self.0.iter().map(|i| *i)
+    }
+}
+
+pub(crate) trait Encoder: Send + Write {
+    fn finish(self: Box<Self>) -> Result<(), Error>;
+}
+
+impl<W: Send + Write> Encoder for GzEncoder<W> {
+    fn finish(self: Box<Self>) -> Result<(), Error> {
+        GzEncoder::finish(*self).context("failed to finish .gz file")?;
+        Ok(())
+    }
+}
+
+impl<W: Send + Write> Encoder for XzEncoder<W> {
+    fn finish(self: Box<Self>) -> Result<(), Error> {
+        XzEncoder::finish(*self).context("failed to finish .xz file")?;
+        Ok(())
+    }
+}
+
+pub(crate) struct CombinedEncoder {
+    encoders: Vec<Box<dyn Encoder>>,
+}
+
+impl CombinedEncoder {
+    pub(crate) fn new(encoders: Vec<Box<dyn Encoder>>) -> Box<dyn Encoder> {
+        Box::new(Self { encoders })
+    }
+}
+
+impl Write for CombinedEncoder {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.write_all(buf)?;
+        Ok(buf.len())
+    }
+
+    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
+        self.encoders
+            .par_iter_mut()
+            .map(|w| w.write_all(buf))
+            .collect::<std::io::Result<Vec<()>>>()?;
+        Ok(())
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        self.encoders
+            .par_iter_mut()
+            .map(|w| w.flush())
+            .collect::<std::io::Result<Vec<()>>>()?;
+        Ok(())
+    }
+}
+
+impl Encoder for CombinedEncoder {
+    fn finish(self: Box<Self>) -> Result<(), Error> {
+        self.encoders
+            .into_par_iter()
+            .map(|e| e.finish())
+            .collect::<Result<Vec<()>, Error>>()?;
+        Ok(())
+    }
+}
diff --git a/src/generator.rs b/src/generator.rs
@@ -1,5 +1,6 @@
 use super::Scripter;
 use super::Tarballer;
+use crate::compression::CompressionFormats;
 use crate::util::*;
 use anyhow::{bail, format_err, Context, Result};
 use std::io::Write;
@@ -40,6 +41,9 @@ actor! {
 
         /// The location to put the final image and tarball
         output_dir: String = "./dist",
+
+        /// The formats used to compress the tarball
+        compression_formats: CompressionFormats = CompressionFormats::default(),
     }
 }
 
@@ -85,7 +89,7 @@ impl Generator {
             .rel_manifest_dir(self.rel_manifest_dir)
             .success_message(self.success_message)
             .legacy_manifest_dirs(self.legacy_manifest_dirs)
-            .output_script(path_to_str(&output_script)?);
+            .output_script(path_to_str(&output_script)?.into());
         scripter.run()?;
 
         // Make the tarballs
@@ -95,7 +99,8 @@ impl Generator {
         tarballer
             .work_dir(self.work_dir)
             .input(self.package_name)
-            .output(path_to_str(&output)?);
+            .output(path_to_str(&output)?.into())
+            .compression_formats(self.compression_formats.clone());
         tarballer.run()?;
 
         Ok(())
diff --git a/src/lib.rs b/src/lib.rs
@@ -2,6 +2,7 @@
 mod util;
 
 mod combiner;
+mod compression;
 mod generator;
 mod scripter;
 mod tarballer;
diff --git a/src/main.rs b/src/main.rs
@@ -1,5 +1,6 @@
 use anyhow::{Context, Result};
 use clap::{App, ArgMatches};
+use std::convert::TryInto;
 
 fn main() -> Result<()> {
     let yaml = clap::load_yaml!("main.yml");
@@ -19,7 +20,11 @@ macro_rules! parse(
     ($matches:expr => $type:ty { $( $option:tt => $setter:ident, )* }) => {
         {
             let mut command: $type = Default::default();
-            $( $matches.value_of($option).map(|s| command.$setter(s)); )*
+            $(
+                if let Some(val) = $matches.value_of($option) {
+                    command.$setter(val.try_into()?);
+                }
+            )*
             command
         }
     }
@@ -36,6 +41,7 @@ fn combine(matches: &ArgMatches<'_>) -> Result<()> {
         "non-installed-overlay" => non_installed_overlay,
         "work-dir" => work_dir,
         "output-dir" => output_dir,
+        "compression-formats" => compression_formats,
     });
 
     combiner.run().context("failed to combine installers")?;
@@ -55,6 +61,7 @@ fn generate(matches: &ArgMatches<'_>) -> Result<()> {
         "image-dir" => image_dir,
         "work-dir" => work_dir,
         "output-dir" => output_dir,
+        "compression-formats" => compression_formats,
     });
 
     generator.run().context("failed to generate installer")?;
@@ -81,6 +88,7 @@ fn tarball(matches: &ArgMatches<'_>) -> Result<()> {
         "input" => input,
         "output" => output,
         "work-dir" => work_dir,
+        "compression-formats" => compression_formats,
     });
 
     tarballer.run().context("failed to generate tarballs")?;
diff --git a/src/main.yml b/src/main.yml
@@ -60,6 +60,11 @@ subcommands:
             long: output-dir
             takes_value: true
             value_name: DIR
+        - compression-formats:
+            help: Comma-separated list of compression formats to use
+            long: compression-formats
+            takes_value: true
+            value_name: FORMAT
   - combine:
       about: Combine installer tarballs
       args:
@@ -108,6 +113,11 @@ subcommands:
             long: output-dir
             takes_value: true
             value_name: DIR
+        - compression-formats:
+            help: Comma-separated list of compression formats to use
+            long: compression-formats
+            takes_value: true
+            value_name: FORMAT
   - script:
       about: Generate an installation script
       args:
@@ -154,4 +164,9 @@ subcommands:
             long: work-dir
             takes_value: true
             value_name: DIR
+        - compression-formats:
+            help: Comma-separated list of compression formats to use
+            long: compression-formats
+            takes_value: true
+            value_name: FORMAT
 
diff --git a/src/tarballer.rs b/src/tarballer.rs
diff --git a/src/util.rs b/src/util.rs
diff --git a/test.sh b/test.sh
