Auto merge of #613 - Mark-Simulacrum:opt, r=Mark-Simulacrum

Reduce memory pressure when writing out log archives

The S3 archive writing was forced to keep two copies of the archive in memory, which is updated to just one archive by this commit. We may wish to incrementally write out the report in the future, but that will take a good bit more work.

Author: bors

src/report/archives.rs

@@ -2,7 +2,7 @@ use crate::config::Config;
 use crate::crates::Crate;
 use crate::experiments::Experiment;
 use crate::prelude::*;
-use crate::report::{compare, ReportWriter};
+use crate::report::{compare, Comparison, ReportWriter};
 use crate::results::{EncodedLog, EncodingType, ReadResults};
 use flate2::{write::GzEncoder, Compression};
 use indexmap::IndexMap;
@@ -14,77 +14,151 @@ pub struct Archive {
     path: String,
 }
 
-pub fn write_logs_archives<DB: ReadResults, W: ReportWriter>(
+struct LogEntry {
+    path: String,
+    comparison: Comparison,
+    log_bytes: Vec<u8>,
+}
+
+impl LogEntry {
+    fn header(&self) -> TarHeader {
+        let mut header = TarHeader::new_gnu();
+        header.set_size(self.log_bytes.len() as u64);
+        header.set_mode(0o644);
+        header.set_cksum();
+        header
+    }
+}
+
+fn iterate<'a, DB: ReadResults + 'a>(
+    db: &'a DB,
+    ex: &'a Experiment,
+    crates: &'a [Crate],
+    config: &'a Config,
+) -> impl Iterator<Item = Fallible<LogEntry>> + 'a {
+    let mut iter = crates
+        .iter()
+        .filter(move |krate| !config.should_skip(krate))
+        .map(move |krate| -> Fallible<Vec<LogEntry>> {
+            let res1 = db.load_test_result(ex, &ex.toolchains[0], krate)?;
+            let res2 = db.load_test_result(ex, &ex.toolchains[1], krate)?;
+            let comparison = compare(config, krate, res1.as_ref(), res2.as_ref());
+
+            ex.toolchains
+                .iter()
+                .filter_map(move |tc| {
+                    let log = db
+                        .load_log(ex, tc, krate)
+                        .and_then(|c| c.ok_or_else(|| err_msg("missing logs")))
+                        .with_context(|_| format!("failed to read log of {} on {}", krate, tc));
+
+                    let log_bytes: EncodedLog = match log {
+                        Ok(l) => l,
+                        Err(e) => {
+                            crate::utils::report_failure(&e);
+                            return None;
+                        }
+                    };
+
+                    let log_bytes = match log_bytes.to_plain() {
+                        Ok(it) => it,
+                        Err(err) => return Some(Err(err)),
+                    };
+
+                    let path = format!(
+                        "{}/{}/{}.txt",
+                        comparison,
+                        krate.id(),
+                        tc.to_path_component(),
+                    );
+                    Some(Ok(LogEntry {
+                        path,
+                        comparison,
+                        log_bytes,
+                    }))
+                })
+                .collect()
+        });
+
+    let mut in_progress = vec![].into_iter();
+    std::iter::from_fn(move || loop {
+        if let Some(next) = in_progress.next() {
+            return Some(Ok(next));
+        }
+        match iter.next()? {
+            Ok(list) => in_progress = list.into_iter(),
+            Err(err) => return Some(Err(err)),
+        }
+    })
+}
+
+fn write_all_archive<DB: ReadResults, W: ReportWriter>(
     db: &DB,
     ex: &Experiment,
     crates: &[Crate],
     dest: &W,
     config: &Config,
-) -> Fallible<Vec<Archive>> {
-    let mut archives = Vec::new();
-    let mut all = TarBuilder::new(GzEncoder::new(Vec::new(), Compression::default()));
-    let mut by_comparison = IndexMap::new();
-
-    for krate in crates {
-        if config.should_skip(krate) {
-            continue;
+) -> Fallible<Archive> {
+    for i in 1..=RETRIES {
+        let mut all = TarBuilder::new(GzEncoder::new(Vec::new(), Compression::default()));
+        for entry in iterate(db, ex, crates, config) {
+            let entry = entry?;
+            let mut header = entry.header();
+            all.append_data(&mut header, &entry.path, &entry.log_bytes[..])?;
         }
 
-        let res1 = db.load_test_result(ex, &ex.toolchains[0], krate)?;
-        let res2 = db.load_test_result(ex, &ex.toolchains[1], krate)?;
-        let comparison = compare(config, krate, res1.as_ref(), res2.as_ref());
-
-        for tc in &ex.toolchains {
-            let log = db
-                .load_log(ex, tc, krate)
-                .and_then(|c| c.ok_or_else(|| err_msg("missing logs")))
-                .with_context(|_| format!("failed to read log of {} on {}", krate, tc));
-
-            let log_bytes: EncodedLog = match log {
-                Ok(l) => l,
-                Err(e) => {
-                    crate::utils::report_failure(&e);
+        let data = all.into_inner()?.finish()?;
+        let len = data.len();
+        match dest.write_bytes_once(
+            "logs-archives/all.tar.gz",
+            data,
+            &"application/gzip".parse().unwrap(),
+            EncodingType::Plain,
+        ) {
+            Ok(()) => break,
+            Err(e) => {
+                if i == RETRIES {
+                    return Err(e);
+                } else {
+                    std::thread::sleep(std::time::Duration::from_secs(2));
+                    warn!(
+                        "retry ({}/{}) writing logs-archives/all.tar.gz ({} bytes)",
+                        i, RETRIES, len,
+                    );
                     continue;
                 }
-            };
-
-            let log_bytes = log_bytes.to_plain()?;
-            let log_bytes = log_bytes.as_slice();
-
-            let path = format!(
-                "{}/{}/{}.txt",
-                comparison,
-                krate.id(),
-                tc.to_path_component(),
-            );
-
-            let mut header = TarHeader::new_gnu();
-            header.set_size(log_bytes.len() as u64);
-            header.set_mode(0o644);
-            header.set_cksum();
-
-            all.append_data(&mut header, &path, log_bytes)?;
-            by_comparison
-                .entry(comparison)
-                .or_insert_with(|| {
-                    TarBuilder::new(GzEncoder::new(Vec::new(), Compression::default()))
-                })
-                .append_data(&mut header, &path, log_bytes)?;
+            }
         }
     }
 
-    let data = all.into_inner()?.finish()?;
-    dest.write_bytes(
-        "logs-archives/all.tar.gz",
-        data,
-        &"application/gzip".parse().unwrap(),
-        EncodingType::Plain,
-    )?;
-
-    archives.push(Archive {
+    Ok(Archive {
         name: "All the crates".to_string(),
         path: "logs-archives/all.tar.gz".to_string(),
-    });
+    })
+}
+
+const RETRIES: usize = 4;
+
+pub fn write_logs_archives<DB: ReadResults, W: ReportWriter>(
+    db: &DB,
+    ex: &Experiment,
+    crates: &[Crate],
+    dest: &W,
+    config: &Config,
+) -> Fallible<Vec<Archive>> {
+    let mut archives = Vec::new();
+    let mut by_comparison = IndexMap::new();
+
+    archives.push(write_all_archive(db, ex, crates, dest, config)?);
+
+    for entry in iterate(db, ex, crates, config) {
+        let entry = entry?;
+
+        by_comparison
+            .entry(entry.comparison)
+            .or_insert_with(|| TarBuilder::new(GzEncoder::new(Vec::new(), Compression::default())))
+            .append_data(&mut entry.header(), &entry.path, &entry.log_bytes[..])?;
+    }
 
     for (comparison, archive) in by_comparison.drain(..) {
         let data = archive.into_inner()?.finish()?;

src/report/mod.rs

@@ -512,6 +512,17 @@ pub trait ReportWriter {
         mime: &Mime,
         encoding_type: EncodingType,
     ) -> Fallible<()>;
+    // This doesn't retry writing -- lets the caller handle failures.
+    //
+    // This is done to avoid copying the bytes buffer, since rusoto at least
+    // currently requires ownership.
+    fn write_bytes_once<P: AsRef<Path>>(
+        &self,
+        path: P,
+        b: Vec<u8>,
+        mime: &Mime,
+        encoding_type: EncodingType,
+    ) -> Fallible<()>;
     fn write_string<P: AsRef<Path>>(&self, path: P, s: Cow<str>, mime: &Mime) -> Fallible<()>;
     fn copy<P: AsRef<Path>, R: Read>(&self, r: &mut R, path: P, mime: &Mime) -> Fallible<()>;
 }
@@ -544,6 +555,16 @@ impl ReportWriter for FileWriter {
         Ok(())
     }
 
+    fn write_bytes_once<P: AsRef<Path>>(
+        &self,
+        path: P,
+        b: Vec<u8>,
+        mime: &Mime,
+        encoding: EncodingType,
+    ) -> Fallible<()> {
+        self.write_bytes(path, b, mime, encoding)
+    }
+
     fn write_string<P: AsRef<Path>>(&self, path: P, s: Cow<str>, _: &Mime) -> Fallible<()> {
         self.create_prefix(path.as_ref())?;
         fs::write(&self.0.join(path.as_ref()), s.as_ref().as_bytes())?;
@@ -595,6 +616,19 @@ impl ReportWriter for DummyWriter {
         Ok(())
     }
 
+    fn write_bytes_once<P: AsRef<Path>>(
+        &self,
+        path: P,
+        b: Vec<u8>,
+        mime: &Mime,
+        _: EncodingType,
+    ) -> Fallible<()> {
+        self.results
+            .borrow_mut()
+            .insert((path.as_ref().to_path_buf(), mime.clone()), b);
+        Ok(())
+    }
+
     fn write_string<P: AsRef<Path>>(&self, path: P, s: Cow<str>, mime: &Mime) -> Fallible<()> {
         self.results.borrow_mut().insert(
             (path.as_ref().to_path_buf(), mime.clone()),

src/report/s3.rs

@@ -144,6 +144,40 @@ impl ReportWriter for S3Writer {
         }
     }
 
+    fn write_bytes_once<P: AsRef<Path>>(
+        &self,
+        path: P,
+        s: Vec<u8>,
+        mime: &Mime,
+        encoding_type: EncodingType,
+    ) -> Fallible<()> {
+        let req = PutObjectRequest {
+            acl: Some("public-read".into()),
+            body: Some(s.into()),
+            bucket: self.prefix.bucket.clone(),
+            key: self
+                .prefix
+                .prefix
+                .join(path.as_ref())
+                .to_string_lossy()
+                .into(),
+            content_type: Some(mime.to_string()),
+            content_encoding: match encoding_type {
+                EncodingType::Plain => None,
+                EncodingType::Gzip => Some("gzip".into()),
+            },
+            ..Default::default()
+        };
+        let r = self.client.put_object(req).sync();
+        if let Err(::rusoto_s3::PutObjectError::Unknown(ref resp)) = r {
+            error!("S3 request status: {}", resp.status);
+            error!("S3 request body: {}", String::from_utf8_lossy(&resp.body));
+            error!("S3 request headers: {:?}", resp.headers);
+        }
+        r.with_context(|_| format!("S3 failure to upload {:?}", path.as_ref()))?;
+        Ok(())
+    }
+
     fn write_string<P: AsRef<Path>>(&self, path: P, s: Cow<str>, mime: &Mime) -> Fallible<()> {
         self.write_bytes(path, s.into_owned().into_bytes(), mime, EncodingType::Plain)
     }