Back the largest archive by a tempfile (backed by disk)

This allows for larger archives to be written without impacting our
memory usage, which typically leads to OOMs and general bad behavior.

Author: Mark-Simulacrum

Cargo.lock

@@ -59,6 +59,7 @@ tokio = "1.24"
 aws-sdk-s3 = "1.7"
 aws-config = { version = "1", features = ["behavior-version-latest"] }
 thiserror = "1.0.38"
+nix = { version = "0.27.1", features = ["mman"] }
 
 [dev-dependencies]
 assert_cmd = "2.0.4"

Cargo.toml

@@ -1,3 +1,7 @@
+use std::fs::File;
+use std::num::NonZeroUsize;
+use std::ptr::NonNull;
+
 use crate::config::Config;
 use crate::crates::Crate;
 use crate::experiments::Experiment;
@@ -7,6 +11,49 @@ use crate::results::{EncodedLog, EncodingType, ReadResults};
 use flate2::{write::GzEncoder, Compression};
 use indexmap::IndexMap;
 use tar::{Builder as TarBuilder, Header as TarHeader};
+use tempfile::tempfile;
+
+struct TempfileBackedBuffer {
+    _file: File,
+    mmap: NonNull<[u8]>,
+}
+
+impl TempfileBackedBuffer {
+    fn new(file: File) -> Fallible<TempfileBackedBuffer> {
+        let len = file.metadata()?.len().try_into().unwrap();
+        unsafe {
+            let base = nix::sys::mman::mmap(
+                None,
+                NonZeroUsize::new(len).unwrap(),
+                nix::sys::mman::ProtFlags::PROT_READ,
+                nix::sys::mman::MapFlags::MAP_PRIVATE,
+                Some(&file),
+                0,
+            )?;
+            let Some(base) = NonNull::new(base as *mut u8) else {
+                panic!("Failed to map file");
+            };
+            Ok(TempfileBackedBuffer {
+                _file: file,
+                mmap: NonNull::slice_from_raw_parts(base, len),
+            })
+        }
+    }
+
+    fn buffer(&self) -> &[u8] {
+        unsafe { self.mmap.as_ref() }
+    }
+}
+
+impl Drop for TempfileBackedBuffer {
+    fn drop(&mut self) {
+        unsafe {
+            if let Err(e) = nix::sys::mman::munmap(self.mmap.as_ptr() as *mut _, self.mmap.len()) {
+                eprintln!("Failed to unmap temporary file: {:?}", e);
+            }
+        }
+    }
+}
 
 #[derive(Serialize)]
 pub struct Archive {
@@ -100,18 +147,23 @@ fn write_all_archive<DB: ReadResults, W: ReportWriter>(
     config: &Config,
 ) -> Fallible<Archive> {
     for i in 1..=RETRIES {
-        let mut all = TarBuilder::new(GzEncoder::new(Vec::new(), Compression::default()));
+        // We write this large-ish tarball into a tempfile, which moves the I/O to disk operations
+        // rather than keeping it in memory. This avoids complicating the code by doing incremental
+        // writes to S3 (requiring buffer management etc) while avoiding keeping the blob entirely
+        // in memory.
+        let backing = tempfile()?;
+        let mut all = TarBuilder::new(GzEncoder::new(backing, 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 data = all.into_inner()?.finish()?;
-        let len = data.len();
+        let buffer = TempfileBackedBuffer::new(data)?;
         match dest.write_bytes(
             "logs-archives/all.tar.gz",
-            &data,
+            buffer.buffer(),
             &"application/gzip".parse().unwrap(),
             EncodingType::Plain,
         ) {
@@ -123,7 +175,10 @@ fn write_all_archive<DB: ReadResults, W: ReportWriter>(
                     std::thread::sleep(std::time::Duration::from_secs(2));
                     warn!(
                         "retry ({}/{}) writing logs-archives/all.tar.gz ({} bytes) (error: {:?})",
-                        i, RETRIES, len, e,
+                        i,
+                        RETRIES,
+                        buffer.buffer().len(),
+                        e,
                     );
                     continue;
                 }

src/report/archives.rs

@@ -626,7 +626,7 @@ impl ReportWriter for DummyWriter {
         self.results
             .lock()
             .unwrap()
-            .insert((path.as_ref().to_path_buf(), mime.clone()), b);
+            .insert((path.as_ref().to_path_buf(), mime.clone()), b.to_vec());
         Ok(())
     }