diff --git a/src/combiner.rs b/src/combiner.rs index 9cdd2df..006a40c 100644 --- 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 new file mode 100644 index 0000000..b3010cb --- /dev/null +++ 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) -> Option { + 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) -> Result, 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) -> Result, 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 in order to parse the value from the command line. +#[derive(Debug, Clone)] +pub struct CompressionFormats(Vec); + +impl TryFrom<&'_ str> for CompressionFormats { + type Error = Error; + + fn try_from(value: &str) -> Result { + 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 + '_ { + self.0.iter().map(|i| *i) + } +} + +pub(crate) trait Encoder: Send + Write { + fn finish(self: Box) -> Result<(), Error>; +} + +impl Encoder for GzEncoder { + fn finish(self: Box) -> Result<(), Error> { + GzEncoder::finish(*self).context("failed to finish .gz file")?; + Ok(()) + } +} + +impl Encoder for XzEncoder { + fn finish(self: Box) -> Result<(), Error> { + XzEncoder::finish(*self).context("failed to finish .xz file")?; + Ok(()) + } +} + +pub(crate) struct CombinedEncoder { + encoders: Vec>, +} + +impl CombinedEncoder { + pub(crate) fn new(encoders: Vec>) -> Box { + Box::new(Self { encoders }) + } +} + +impl Write for CombinedEncoder { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + 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::>>()?; + Ok(()) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.encoders + .par_iter_mut() + .map(|w| w.flush()) + .collect::>>()?; + Ok(()) + } +} + +impl Encoder for CombinedEncoder { + fn finish(self: Box) -> Result<(), Error> { + self.encoders + .into_par_iter() + .map(|e| e.finish()) + .collect::, Error>>()?; + Ok(()) + } +} diff --git a/src/generator.rs b/src/generator.rs index ecf73ca..2601eb5 100644 --- 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 index 8e88311..7990920 100644 --- 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 index 8b5f1e1..e933dd0 100644 --- 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 index 2f9978b..7b6d735 100644 --- 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 index 42a4ffa..4ac8cf7 100644 --- a/src/tarballer.rs +++ b/src/tarballer.rs @@ -1,13 +1,14 @@ use anyhow::{bail, Context, Result}; -use flate2::write::GzEncoder; use std::fs::{read_link, symlink_metadata}; -use std::io::{self, empty, BufWriter, Write}; +use std::io::{empty, BufWriter, Write}; use std::path::Path; use tar::{Builder, Header}; use walkdir::WalkDir; -use xz2::write::XzEncoder; -use crate::util::*; +use crate::{ + compression::{CombinedEncoder, CompressionFormats}, + util::*, +}; actor! { #[derive(Debug)] @@ -20,21 +21,22 @@ actor! { /// The folder in which the input is to be found. work_dir: String = "./workdir", + + /// The formats used to compress the tarball. + compression_formats: CompressionFormats = CompressionFormats::default(), } } impl Tarballer { /// Generates the actual tarballs pub fn run(self) -> Result<()> { - let tar_gz = self.output.clone() + ".tar.gz"; - let tar_xz = self.output.clone() + ".tar.xz"; - - // Remove any existing files. - for file in &[&tar_gz, &tar_xz] { - if Path::new(file).exists() { - remove_file(file)?; - } - } + let tarball_name = self.output.clone() + ".tar"; + let encoder = CombinedEncoder::new( + self.compression_formats + .iter() + .map(|f| f.encode(&tarball_name)) + .collect::>>()?, + ); // Sort files by their suffix, to group files with the same name from // different locations (likely identical) and files with the same @@ -43,22 +45,9 @@ impl Tarballer { .context("failed to collect file paths")?; files.sort_by(|a, b| a.bytes().rev().cmp(b.bytes().rev())); - // Prepare the `.tar.gz` file. - let gz = GzEncoder::new(create_new_file(tar_gz)?, flate2::Compression::best()); - - // Prepare the `.tar.xz` file. 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()?; - let xz = XzEncoder::new_stream(create_new_file(tar_xz)?, stream); - // Write the tar into both encoded files. We write all directories // first, so files may be directly created. (See rust-lang/rustup.rs#1092.) - let tee = RayonTee(xz, gz); - let buf = BufWriter::with_capacity(1024 * 1024, tee); + let buf = BufWriter::with_capacity(1024 * 1024, encoder); let mut builder = Builder::new(buf); let pool = rayon::ThreadPoolBuilder::new() @@ -77,20 +66,14 @@ impl Tarballer { append_path(&mut builder, &src, &path) .with_context(|| format!("failed to tar file '{}'", src.display()))?; } - let RayonTee(xz, gz) = builder + builder .into_inner() .context("failed to finish writing .tar stream")? .into_inner() .ok() - .unwrap(); + .unwrap() + .finish()?; - // Finish both encoded files. - let (rxz, rgz) = rayon::join( - || xz.finish().context("failed to finish .tar.xz file"), - || gz.finish().context("failed to finish .tar.gz file"), - ); - rxz?; - rgz?; Ok(()) }) } @@ -154,24 +137,3 @@ where } Ok((dirs, files)) } - -struct RayonTee(A, B); - -impl Write for RayonTee { - fn write(&mut self, buf: &[u8]) -> io::Result { - self.write_all(buf)?; - Ok(buf.len()) - } - - fn write_all(&mut self, buf: &[u8]) -> io::Result<()> { - let (a, b) = (&mut self.0, &mut self.1); - let (ra, rb) = rayon::join(|| a.write_all(buf), || b.write_all(buf)); - ra.and(rb) - } - - fn flush(&mut self) -> io::Result<()> { - let (a, b) = (&mut self.0, &mut self.1); - let (ra, rb) = rayon::join(|| a.flush(), || b.flush()); - ra.and(rb) - } -} diff --git a/src/util.rs b/src/util.rs index 3ddcfc6..078ceb3 100644 --- a/src/util.rs +++ b/src/util.rs @@ -142,8 +142,8 @@ macro_rules! actor { impl $name { $( $( #[ $field_attr ] )+ - pub fn $field>(&mut self, value: T) -> &mut Self { - self.$field = value.into(); + pub fn $field(&mut self, value: $type) -> &mut Self { + self.$field = value; self })+ } diff --git a/test.sh b/test.sh index dc7275b..bf6de4c 100755 --- a/test.sh +++ b/test.sh @@ -1164,6 +1164,179 @@ docdir_combined() { } runtest docdir_combined +combine_installers_different_input_compression_formats() { + try sh "$S/gen-installer.sh" \ + --image-dir="$TEST_DIR/image1" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name=rustc \ + --component-name=rustc \ + --compression-formats=xz + try sh "$S/gen-installer.sh" \ + --image-dir="$TEST_DIR/image3" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name=cargo \ + --component-name=cargo \ + --compression-formats=gz + try sh "$S/combine-installers.sh" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name=rust \ + --input-tarballs="$OUT_DIR/rustc.tar.xz,$OUT_DIR/cargo.tar.gz" + + try test -e "${OUT_DIR}/rust.tar.gz" + try test -e "${OUT_DIR}/rust.tar.xz" +} +runtest combine_installers_different_input_compression_formats + +generate_compression_formats_one() { + try sh "$S/gen-installer.sh" \ + --image-dir="$TEST_DIR/image1" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name="rustc" \ + --component-name="rustc" \ + --compression-formats="xz" + + try test ! -e "${OUT_DIR}/rustc.tar.gz" + try test -e "${OUT_DIR}/rustc.tar.xz" +} +runtest generate_compression_formats_one + +generate_compression_formats_multiple() { + try sh "$S/gen-installer.sh" \ + --image-dir="$TEST_DIR/image1" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name="rustc" \ + --component-name="rustc" \ + --compression-formats="gz,xz" + + try test -e "${OUT_DIR}/rustc.tar.gz" + try test -e "${OUT_DIR}/rustc.tar.xz" +} +runtest generate_compression_formats_multiple + +generate_compression_formats_error() { + expect_fail sh "$S/gen-installer.sh" \ + --image-dir="$TEST_DIR/image1" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name="rustc" \ + --component-name="rustc" \ + --compression-formats="xz,foobar" +} +runtest generate_compression_formats_error + +combine_compression_formats_one() { + try sh "$S/gen-installer.sh" \ + --image-dir="$TEST_DIR/image1" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name=rustc \ + --component-name=rustc + try sh "$S/gen-installer.sh" \ + --image-dir="$TEST_DIR/image3" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name=cargo \ + --component-name=cargo + try sh "$S/combine-installers.sh" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name=rust \ + --input-tarballs="$OUT_DIR/rustc.tar.gz,$OUT_DIR/cargo.tar.gz" \ + --compression-formats=xz + + try test ! -e "${OUT_DIR}/rust.tar.gz" + try test -e "${OUT_DIR}/rust.tar.xz" +} +runtest combine_compression_formats_one + +combine_compression_formats_multiple() { + try sh "$S/gen-installer.sh" \ + --image-dir="$TEST_DIR/image1" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name=rustc \ + --component-name=rustc + try sh "$S/gen-installer.sh" \ + --image-dir="$TEST_DIR/image3" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name=cargo \ + --component-name=cargo + try sh "$S/combine-installers.sh" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name=rust \ + --input-tarballs="$OUT_DIR/rustc.tar.gz,$OUT_DIR/cargo.tar.gz" \ + --compression-formats=xz,gz + + try test -e "${OUT_DIR}/rust.tar.gz" + try test -e "${OUT_DIR}/rust.tar.xz" +} +runtest combine_compression_formats_multiple + +combine_compression_formats_error() { + try sh "$S/gen-installer.sh" \ + --image-dir="$TEST_DIR/image1" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name=rustc \ + --component-name=rustc + try sh "$S/gen-installer.sh" \ + --image-dir="$TEST_DIR/image3" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name=cargo \ + --component-name=cargo + expect_fail sh "$S/combine-installers.sh" \ + --work-dir="$WORK_DIR" \ + --output-dir="$OUT_DIR" \ + --package-name=rust \ + --input-tarballs="$OUT_DIR/rustc.tar.gz,$OUT_DIR/cargo.tar.gz" \ + --compression-formats=xz,foobar +} +runtest combine_compression_formats_error + +tarball_compression_formats_one() { + try cp -r "${TEST_DIR}/image1" "${WORK_DIR}/image" + try sh "$S/make-tarballs.sh" \ + --input="${WORK_DIR}/image" \ + --work-dir="${WORK_DIR}" \ + --output="${OUT_DIR}/rustc" \ + --compression-formats="xz" + + try test ! -e "${OUT_DIR}/rustc.tar.gz" + try test -e "${OUT_DIR}/rustc.tar.xz" +} +runtest tarball_compression_formats_one + +tarball_compression_formats_multiple() { + try cp -r "${TEST_DIR}/image1" "${WORK_DIR}/image" + try sh "$S/make-tarballs.sh" \ + --input="${WORK_DIR}/image" \ + --work-dir="${WORK_DIR}" \ + --output="${OUT_DIR}/rustc" \ + --compression-formats="xz,gz" + + try test -e "${OUT_DIR}/rustc.tar.gz" + try test -e "${OUT_DIR}/rustc.tar.xz" +} +runtest tarball_compression_formats_multiple + +tarball_compression_formats_error() { + try cp -r "${TEST_DIR}/image1" "${WORK_DIR}/image" + expect_fail sh "$S/make-tarballs.sh" \ + --input="${WORK_DIR}/image" \ + --work-dir="${WORK_DIR}" \ + --output="${OUT_DIR}/rustc" \ + --compression-formats="xz,foobar" +} +runtest tarball_compression_formats_error + echo echo "TOTAL SUCCESS!" echo