Merge pull request #76 from cuviper/rayon

Stream the parallel xz/gz tarball generation

Author: alexcrichton

Cargo.toml

@@ -11,9 +11,10 @@ path = "src/main.rs"
 [dependencies]
 error-chain = "0.11.0"
 flate2 = "1.0.1"
+rayon = "0.9"
 tar = "0.4.13"
 walkdir = "1.0.7"
-xz2 = "0.1.3"
+xz2 = "0.1.4"
 
 [dependencies.clap]
 features = ["yaml"]

src/lib.rs

@@ -11,6 +11,7 @@
 #[macro_use]
 extern crate error_chain;
 extern crate flate2;
+extern crate rayon;
 extern crate tar;
 extern crate walkdir;
 extern crate xz2;

src/tarballer.rs

@@ -9,13 +9,12 @@
 // except according to those terms.
 
 use std::fs::File;
-use std::io::Write;
+use std::io::{self, Write, BufWriter};
 use std::path::Path;
-use std::sync::Arc;
-use std::thread;
 
 use flate2;
 use flate2::write::GzEncoder;
+use rayon;
 use tar::{Builder, Header};
 use walkdir::WalkDir;
 use xz2::write::XzEncoder;
@@ -57,41 +56,42 @@ impl Tarballer {
             .chain_err(|| "failed to collect file paths")?;
         files.sort_by(|a, b| a.bytes().rev().cmp(b.bytes().rev()));
 
-        // Write the tar into both encoded files.  We write all directories
-        // first, so files may be directly created. (see rustup.rs#1092)
-        let mut builder = Builder::new(Vec::new());
-        for path in dirs {
-            let src = Path::new(&self.work_dir).join(&path);
-            builder.append_dir(&path, &src)
-                .chain_err(|| format!("failed to tar dir '{}'", src.display()))?;
-        }
-        for path in files {
-            let src = Path::new(&self.work_dir).join(&path);
-            let file = open_file(&src)?;
-            builder.append_data(&mut header(&src, &file)?, &path, &file)
-                .chain_err(|| format!("failed to tar file '{}'", src.display()))?;
-        }
-        let contents = builder.into_inner()
-            .chain_err(|| "failed to finish writing .tar stream")?;
-        let contents = Arc::new(contents);
-
         // Prepare the .tar.gz file
-        let contents2 = contents.clone();
-        let t = thread::spawn(move || {
-            let mut gz = GzEncoder::new(create_new_file(tar_gz)?,
-                                        flate2::Compression::best());
-            gz.write_all(&contents2).chain_err(|| "failed to write .gz")?;
-            gz.finish().chain_err(|| "failed to finish .gz")
-        });
+        let gz = GzEncoder::new(create_new_file(tar_gz)?, flate2::Compression::best());
 
         // Prepare the .tar.xz file
-        let mut xz = XzEncoder::new(create_new_file(tar_xz)?, 9);
-        xz.write_all(&contents).chain_err(|| "failed to write .xz")?;
-        xz.finish().chain_err(|| "failed to finish .xz")?;
+        let xz = XzEncoder::new(create_new_file(tar_xz)?, 9);
 
-        t.join().unwrap()?;
-
-        Ok(())
+        // Write the tar into both encoded files.  We write all directories
+        // first, so files may be directly created. (see rustup.rs#1092)
+        let tee = RayonTee(xz, gz);
+        let buf = BufWriter::with_capacity(1024 * 1024, tee);
+        let mut builder = Builder::new(buf);
+
+        let pool = rayon::Configuration::new().num_threads(2).build().unwrap();
+        pool.install(move || {
+            for path in dirs {
+                let src = Path::new(&self.work_dir).join(&path);
+                builder.append_dir(&path, &src)
+                    .chain_err(|| format!("failed to tar dir '{}'", src.display()))?;
+            }
+            for path in files {
+                let src = Path::new(&self.work_dir).join(&path);
+                let file = open_file(&src)?;
+                builder.append_data(&mut header(&src, &file)?, &path, &file)
+                    .chain_err(|| format!("failed to tar file '{}'", src.display()))?;
+            }
+            let RayonTee(xz, gz) = builder.into_inner()
+                .chain_err(|| "failed to finish writing .tar stream")?
+                .into_inner().ok().unwrap();
+
+            // Finish both encoded files
+            let (rxz, rgz) = rayon::join(
+                || xz.finish().chain_err(|| "failed to finish .tar.xz file"),
+                || gz.finish().chain_err(|| "failed to finish .tar.gz file"),
+            );
+            rxz.and(rgz).and(Ok(()))
+        })
     }
 }
 
@@ -138,3 +138,24 @@ fn get_recursive_paths<P, Q>(root: P, name: Q) -> Result<(Vec<String>, Vec<Strin
     }
     Ok((dirs, files))
 }
+
+struct RayonTee<A, B>(A, B);
+
+impl<A: Write + Send, B: Write + Send> Write for RayonTee<A, B> {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        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)
+    }
+}