Compress inputs and outputs over the wire

Author: aidanhs

src/bin/sccache-dist/build.rs

@@ -18,12 +18,12 @@ use libmount::Overlay;
 use lru_disk_cache::Error as LruError;
 use nix;
 use sccache::dist::{
-    BuildResult, CompileCommand, InputsReader, TcCache, Toolchain,
+    BuildResult, CompileCommand, InputsReader, OutputData, TcCache, Toolchain,
     BuilderIncoming,
 };
 use std::collections::HashMap;
 use std::fs;
-use std::io::{self, Read};
+use std::io;
 use std::iter;
 use std::path::{self, Path, PathBuf};
 use std::process::{Command, Output, Stdio};
@@ -208,8 +208,7 @@ impl OverlayBuilder {
                 let abspath = join_suffix(&target_dir, cwd.join(&path)); // Resolve in case it's relative since we copy it from the root level
                 match fs::File::open(abspath) {
                     Ok(mut file) => {
-                        let mut output = vec![];
-                        file.read_to_end(&mut output).unwrap();
+                        let output = OutputData::from_reader(file);
                         outputs.push((path, output))
                     },
                     Err(e) => {
@@ -508,7 +507,7 @@ impl DockerBuilder {
             // TODO: this isn't great, but cp gives it out as a tar
             let output = Command::new("docker").args(&["exec", cid, "/busybox", "cat"]).arg(abspath).output().unwrap();
             if output.status.success() {
-                outputs.push((path, output.stdout))
+                outputs.push((path, OutputData::from_reader(&*output.stdout)))
             } else {
                 debug!("Missing output path {:?}", path)
             }

src/compiler/c.rs

@@ -298,7 +298,7 @@ impl<I: CCompilerImpl> Compilation for CCompilation<I> {
     }
 
     #[cfg(feature = "dist-client")]
-    fn into_dist_inputs_creator(self: Box<Self>, path_transformer: &mut dist::PathTransformer) -> Result<(Box<FnMut(&mut io::Write)>, Box<pkg::CompilerPackager>)> {
+    fn into_dist_inputs_creator(self: Box<Self>, path_transformer: &mut dist::PathTransformer) -> Result<(Box<FnMut(&mut io::Write) + Send>, Box<pkg::CompilerPackager>)> {
         let CCompilation { parsed_args, cwd, preprocessed_input, executable, .. } = *{self};
         trace!("Dist inputs: {:?}", parsed_args.input);
 

src/compiler/compiler.rs

@@ -338,6 +338,7 @@ fn dist_or_local_compile<T>(dist_client: Arc<dist::Client>,
         where T: CommandCreatorSync {
     use boxfnonce::BoxFnOnce;
     use futures::future;
+    use std::io;
 
     debug!("[{}]: Attempting distributed compilation", out_pretty);
     let compile_out_pretty = out_pretty.clone();
@@ -399,9 +400,12 @@ fn dist_or_local_compile<T>(dist_client: Arc<dist::Client>,
                         dist::RunJobResult::Complete(jc) => jc,
                         dist::RunJobResult::JobNotFound => panic!(),
                     };
-                    info!("fetched {:?}", jc.outputs.iter().map(|&(ref p, ref bs)| (p, bs.len())).collect::<Vec<_>>());
-                    for (path, bytes) in jc.outputs {
-                        File::create(path_transformer.to_local(&path)).unwrap().write_all(&bytes).unwrap();
+                    info!("fetched {:?}", jc.outputs.iter().map(|&(ref p, ref bs)| (p, bs.lens().to_string())).collect::<Vec<_>>());
+                    for (path, output_data) in jc.outputs {
+                        let len = output_data.lens().actual;
+                        let mut file = File::create(path_transformer.to_local(&path)).unwrap();
+                        let count = io::copy(&mut output_data.into_reader(), &mut file).unwrap();
+                        assert!(count == len);
                     }
                     jc.output.into()
                 })
@@ -432,7 +436,7 @@ pub trait Compilation {
     // TODO: It's more correct to have a FnBox or Box<FnOnce> here
     #[cfg(feature = "dist-client")]
     fn into_dist_inputs_creator(self: Box<Self>, _path_transformer: &mut dist::PathTransformer)
-                                -> Result<(Box<FnMut(&mut Write)>, Box<CompilerPackager>)> {
+                                -> Result<(Box<FnMut(&mut Write) + Send>, Box<CompilerPackager>)> {
 
         bail!("distributed compilation not implemented")
     }

src/compiler/rust.rs

@@ -1011,7 +1011,7 @@ impl Compilation for RustCompilation {
     }
 
     #[cfg(feature = "dist-client")]
-    fn into_dist_inputs_creator(self: Box<Self>, path_transformer: &mut dist::PathTransformer) -> Result<(Box<FnMut(&mut io::Write)>, Box<pkg::CompilerPackager>)> {
+    fn into_dist_inputs_creator(self: Box<Self>, path_transformer: &mut dist::PathTransformer) -> Result<(Box<FnMut(&mut io::Write) + Send>, Box<pkg::CompilerPackager>)> {
         let RustCompilation { inputs, crate_link_paths, sysroot, crate_types, .. } = *{self};
         trace!("Dist inputs: inputs={:?} crate_link_paths={:?}", inputs, crate_link_paths);
 

src/dist/http.rs

@@ -141,6 +141,7 @@ mod common {
 mod server {
     use bincode;
     use byteorder::{BigEndian, ReadBytesExt};
+    use flate2::read::ZlibDecoder as ZlibReadDecoder;
     use jwt;
     use num_cpus;
     use rand::{self, RngCore};
@@ -526,7 +527,7 @@ mod server {
                         trace!("Req {}: run_job({}): {:?}", req_id, job_id, runjob);
                         let RunJobHttpRequest { command, outputs } = runjob;
                         let body = bincode_reader.into_inner();
-                        let inputs_rdr = InputsReader(Box::new(body));
+                        let inputs_rdr = InputsReader(Box::new(ZlibReadDecoder::new(body)));
                         let outputs = outputs.into_iter().collect();
 
                         let res: RunJobResult = try_or_500_log!(req_id, handler.handle_run_job(&requester, job_id, command, outputs, inputs_rdr));
@@ -566,7 +567,9 @@ mod client {
     use boxfnonce::BoxFnOnce;
     use byteorder::{BigEndian, WriteBytesExt};
     use config;
-    use futures::Future;
+    use flate2::Compression;
+    use flate2::write::ZlibEncoder as ZlibWriteEncoder;
+    use futures::{Future, Stream};
     use futures_cpupool::CpuPool;
     use reqwest;
     use std::fs;
@@ -633,29 +636,37 @@ mod client {
             if let Some(toolchain_file) = self.tc_cache.get_toolchain(&tc) {
                 let url = format!("http://{}/api/v1/distserver/submit_toolchain/{}", job_alloc.server_id.addr(), job_alloc.job_id);
                 let mut req = self.client.post(&url);
-                req.bearer_auth(job_alloc.auth.clone()).body(toolchain_file);
 
-                Box::new(self.pool.spawn_fn(move || bincode_req(&mut req)))
+                Box::new(self.pool.spawn_fn(move || {
+                    req.bearer_auth(job_alloc.auth.clone()).body(toolchain_file);
+                    bincode_req(&mut req)
+                }))
             } else {
                 f_err("couldn't find toolchain locally")
             }
         }
-        fn do_run_job(&self, job_alloc: JobAlloc, command: CompileCommand, outputs: Vec<String>, mut write_inputs: Box<FnMut(&mut Write)>) -> SFuture<RunJobResult> {
+        fn do_run_job(&self, job_alloc: JobAlloc, command: CompileCommand, outputs: Vec<String>, mut write_inputs: Box<FnMut(&mut Write) + Send>) -> SFuture<RunJobResult> {
             let url = format!("http://{}/api/v1/distserver/run_job/{}", job_alloc.server_id.addr(), job_alloc.job_id);
-            let bincode = bincode::serialize(&RunJobHttpRequest { command, outputs }).unwrap();
-            let bincode_length = bincode.len();
-            let mut inputs = vec![];
-            write_inputs(&mut inputs);
+            let mut req = self.client.post(&url);
 
-            let mut body = vec![];
-            body.write_u32::<BigEndian>(bincode_length as u32).unwrap();
-            body.write(&bincode).unwrap();
-            body.write(&inputs).unwrap();
+            Box::new(self.pool.spawn_fn(move || {
+                let bincode = bincode::serialize(&RunJobHttpRequest { command, outputs }).unwrap();
+                let bincode_length = bincode.len();
+
+                let mut body = vec![];
+                body.write_u32::<BigEndian>(bincode_length as u32).unwrap();
+                body.write(&bincode).unwrap();
+                {
+                    let mut compressor = ZlibWriteEncoder::new(&mut body, Compression::fast());
+                    write_inputs(&mut compressor);
+                    compressor.flush().unwrap();
+                    trace!("Compressed inputs from {} -> {}", compressor.total_in(), compressor.total_out());
+                    compressor.finish().unwrap();
+                }
 
-            let mut req = self.client.post(&url);
-            req.bearer_auth(job_alloc.auth.clone()).bytes(body);
-            trace!("Writing {} bytes of input archive", inputs.len());
-            Box::new(self.pool.spawn_fn(move || bincode_req(&mut req)))
+                req.bearer_auth(job_alloc.auth.clone()).bytes(body);
+                bincode_req(&mut req)
+            }))
         }
 
         fn put_toolchain(&self, compiler_path: &Path, weak_key: &str, create: BoxFnOnce<(fs::File,), Result<()>>) -> Result<(Toolchain, Option<String>)> {

src/dist/mod.rs

@@ -16,7 +16,7 @@ use boxfnonce::BoxFnOnce;
 use compiler;
 use std::fmt;
 use std::fs;
-use std::io::{self, Read, Write};
+use std::io::{self, Cursor, Read, Write};
 use std::net::SocketAddr;
 use std::ffi::OsString;
 use std::path::Path;
@@ -271,6 +271,39 @@ impl From<ProcessOutput> for process::Output {
     }
 }
 
+#[derive(Clone, Serialize, Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct OutputData(Vec<u8>, u64);
+impl OutputData {
+    #[cfg(feature = "dist-server")]
+    pub fn from_reader<R: Read>(r: R) -> Self {
+        use flate2::Compression;
+        use flate2::read::ZlibEncoder as ZlibReadEncoder;
+        let mut compressor = ZlibReadEncoder::new(r, Compression::fast());
+        let mut res = vec![];
+        io::copy(&mut compressor, &mut res).unwrap();
+        OutputData(res, compressor.total_in())
+    }
+    pub fn lens(&self) -> OutputDataLens {
+        OutputDataLens { actual: self.1, compressed: self.0.len() as u64 }
+    }
+    #[cfg(feature = "dist-client")]
+    pub fn into_reader(self) -> impl Read {
+        use flate2::read::ZlibDecoder as ZlibReadDecoder;
+        let decompressor = ZlibReadDecoder::new(Cursor::new(self.0));
+        decompressor
+    }
+}
+pub struct OutputDataLens {
+    pub actual: u64,
+    pub compressed: u64,
+}
+impl fmt::Display for OutputDataLens {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "Size: {}->{}", self.actual, self.compressed)
+    }
+}
+
 // TODO: standardise on compressed or not for inputs and toolchain
 
 // TODO: make fields not public
@@ -328,7 +361,7 @@ pub enum RunJobResult {
 #[serde(deny_unknown_fields)]
 pub struct JobComplete {
     pub output: ProcessOutput,
-    pub outputs: Vec<(String, Vec<u8>)>,
+    pub outputs: Vec<(String, OutputData)>,
 }
 
 // Status
@@ -355,7 +388,7 @@ pub enum SubmitToolchainResult {
 
 pub struct BuildResult {
     pub output: ProcessOutput,
-    pub outputs: Vec<(String, Vec<u8>)>,
+    pub outputs: Vec<(String, OutputData)>,
 }
 
 ///////////////////
@@ -431,7 +464,7 @@ pub trait Client {
     // To Server
     // TODO: ideally Box<FnOnce or FnBox
     // BoxFnOnce library doesn't work due to incorrect lifetime inference - https://github.com/rust-lang/rust/issues/28796#issuecomment-410071058
-    fn do_run_job(&self, job_alloc: JobAlloc, command: CompileCommand, outputs: Vec<String>, write_inputs: Box<FnMut(&mut Write)>) -> SFuture<RunJobResult>;
+    fn do_run_job(&self, job_alloc: JobAlloc, command: CompileCommand, outputs: Vec<String>, write_inputs: Box<FnMut(&mut Write) + Send>) -> SFuture<RunJobResult>;
     fn put_toolchain(&self, compiler_path: &Path, weak_key: &str, create: BoxFnOnce<(fs::File,), Result<()>>) -> Result<(Toolchain, Option<String>)>;
     fn may_dist(&self) -> bool;
 }
@@ -447,7 +480,7 @@ impl Client for NoopClient {
     fn do_submit_toolchain(&self, _job_alloc: JobAlloc, _tc: Toolchain) -> SFuture<SubmitToolchainResult> {
         panic!("NoopClient");
     }
-    fn do_run_job(&self, _job_alloc: JobAlloc, _command: CompileCommand, _outputs: Vec<String>, _write_inputs: Box<FnMut(&mut Write)>) -> SFuture<RunJobResult> {
+    fn do_run_job(&self, _job_alloc: JobAlloc, _command: CompileCommand, _outputs: Vec<String>, _write_inputs: Box<FnMut(&mut Write) + Send>) -> SFuture<RunJobResult> {
         panic!("NoopClient");
     }