Skip to content

Commit

Permalink
Merge pull request #105 from rust-lang/dynamic-compression
Browse files Browse the repository at this point in the history
Allow specifying the wanted compression formats at runtime
  • Loading branch information
pietroalbini committed Dec 10, 2020
2 parents 27fed98 + 512629a commit 5254dbf
Show file tree
Hide file tree
Showing 9 changed files with 405 additions and 75 deletions.
38 changes: 25 additions & 13 deletions src/combiner.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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(),
}
}

Expand All @@ -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);

Expand Down Expand Up @@ -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.
Expand All @@ -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(())
Expand Down
154 changes: 154 additions & 0 deletions src/compression.rs
Original file line number Diff line number Diff line change
@@ -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(())
}
}
9 changes: 7 additions & 2 deletions src/generator.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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(),
}
}

Expand Down Expand Up @@ -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
Expand All @@ -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(())
Expand Down
1 change: 1 addition & 0 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
mod util;

mod combiner;
mod compression;
mod generator;
mod scripter;
mod tarballer;
Expand Down
10 changes: 9 additions & 1 deletion src/main.rs
Original file line number Diff line number Diff line change
@@ -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");
Expand All @@ -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
}
}
Expand All @@ -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")?;
Expand All @@ -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")?;
Expand All @@ -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")?;
Expand Down
15 changes: 15 additions & 0 deletions src/main.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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

0 comments on commit 5254dbf

Please sign in to comment.