Extract Cargo artifact iteration into a separate module

Author: Kobzol

collector/src/cargo.rs

@@ -0,0 +1,82 @@
+use anyhow::Context;
+use std::io::BufReader;
+use std::process::{Child, ChildStdout, Command, Stdio};
+
+use cargo_metadata::{Message, MessageIter};
+
+/// Iterator that returns built artifacts from a Cargo command invocation.
+/// It also prints any text lines or messages produced during compilation,
+/// and gathers the messages for better error messages.
+pub struct CargoArtifactIter {
+    stream: MessageIter<BufReader<ChildStdout>>,
+    cargo_process: Child,
+    messages: Vec<String>,
+}
+
+impl CargoArtifactIter {
+    /// Adds arguments to the command required for JSON message parsing, and starts the Cargo
+    /// invocation.
+    pub fn from_cargo_cmd(mut cmd: Command) -> anyhow::Result<Self> {
+        cmd.arg("--message-format")
+            .arg("json-diagnostic-short")
+            .stdin(Stdio::null())
+            .stdout(Stdio::piped())
+            .stderr(Stdio::null());
+        let mut cargo_process = cmd.spawn()?;
+        let stream = BufReader::new(cargo_process.stdout.take().unwrap());
+        Ok(Self {
+            stream: Message::parse_stream(stream),
+            cargo_process,
+            messages: Default::default(),
+        })
+    }
+
+    pub fn finish(mut self) -> anyhow::Result<()> {
+        let output = self
+            .cargo_process
+            .wait()
+            .context("Cargo did not exit successfully")?;
+        if !output.success() {
+            return Err(anyhow::anyhow!(
+                "Failed to run cargo, exit code {}\n{}",
+                output.code().unwrap_or(1),
+                self.messages.join("")
+            ));
+        }
+        Ok(())
+    }
+}
+
+impl Drop for CargoArtifactIter {
+    fn drop(&mut self) {
+        self.cargo_process
+            .wait()
+            .expect("Cargo process did not exit successfully");
+    }
+}
+
+impl Iterator for CargoArtifactIter {
+    type Item = anyhow::Result<cargo_metadata::Artifact>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            match self.stream.next()? {
+                Ok(message) => match message {
+                    Message::CompilerArtifact(artifact) => {
+                        return Some(Ok(artifact));
+                    }
+                    Message::TextLine(line) => {
+                        println!("{line}")
+                    }
+                    Message::CompilerMessage(msg) => {
+                        let message = msg.message.rendered.unwrap_or(msg.message.message);
+                        print!("{message}");
+                        self.messages.push(message);
+                    }
+                    _ => {}
+                },
+                Err(error) => return Some(Err(error.into())),
+            }
+        }
+    }
+}

collector/src/lib.rs

@@ -6,6 +6,7 @@ use std::fmt;
 use std::process::{self, Command};
 
 pub mod api;
+pub mod cargo;
 pub mod codegen;
 pub mod compile;
 pub mod runtime;

collector/src/runtime/benchmark.rs

@@ -1,18 +1,20 @@
-use crate::runtime_group_step_name;
-use crate::toolchain::Toolchain;
-use crate::utils::fs::EnsureImmutableFile;
-use anyhow::Context;
-use benchlib::benchmark::passes_filter;
-use cargo_metadata::Message;
 use core::option::Option;
 use core::option::Option::Some;
 use core::result::Result::Ok;
 use std::collections::HashMap;
-use std::io::BufReader;
 use std::path::{Path, PathBuf};
-use std::process::{Child, Command, Stdio};
+use std::process::Command;
+
+use anyhow::Context;
 use tempfile::TempDir;
 
+use benchlib::benchmark::passes_filter;
+
+use crate::cargo::CargoArtifactIter;
+use crate::runtime_group_step_name;
+use crate::toolchain::Toolchain;
+use crate::utils::fs::EnsureImmutableFile;
+
 /// Directory containing runtime benchmarks.
 /// We measure how long does it take to execute these crates, which is a proxy of the quality
 /// of code generated by rustc.
@@ -217,16 +219,16 @@ pub fn prepare_runtime_benchmark_suite(
             .with_context(|| {
                 anyhow::anyhow!("Cannot start compilation of {}", benchmark_crate.name)
             })
-            .and_then(|process| {
-                parse_benchmark_group(process, &benchmark_crate.name).with_context(|| {
+            .and_then(|iter| {
+                parse_benchmark_group(iter, &benchmark_crate.name).with_context(|| {
                     anyhow::anyhow!("Cannot compile runtime benchmark {}", benchmark_crate.name)
                 })
             });
         match result {
             Ok(group) => groups.push(group),
             Err(error) => {
                 log::error!(
-                    "Cannot compile runtime benchmark group `{}`",
+                    "Cannot compile runtime benchmark group `{}` {error:?}",
                     benchmark_crate.name
                 );
                 failed_to_compile.insert(
@@ -276,66 +278,47 @@ fn check_duplicates(groups: &[BenchmarkGroup]) -> anyhow::Result<()> {
 /// Locates the benchmark binary of a runtime benchmark crate compiled by cargo, and then executes it
 /// to find out what benchmarks do they contain.
 fn parse_benchmark_group(
-    mut cargo_process: Child,
+    mut cargo_iter: CargoArtifactIter,
     group_name: &str,
 ) -> anyhow::Result<BenchmarkGroup> {
     let mut group: Option<BenchmarkGroup> = None;
 
-    let stream = BufReader::new(cargo_process.stdout.take().unwrap());
-    let mut messages = String::new();
-    for message in Message::parse_stream(stream) {
-        let message = message?;
-        match message {
-            Message::CompilerArtifact(artifact) => {
-                if let Some(ref executable) = artifact.executable {
-                    // Found a binary compiled by a runtime benchmark crate.
-                    // Execute it so that we find all the benchmarks it contains.
-                    if artifact.target.kind.iter().any(|k| k == "bin") {
-                        if group.is_some() {
-                            return Err(anyhow::anyhow!("Runtime benchmark group `{group_name}` has produced multiple binaries"));
-                        }
-
-                        let path = executable.as_std_path().to_path_buf();
-                        let benchmarks = gather_benchmarks(&path).map_err(|err| {
-                            anyhow::anyhow!(
-                                "Cannot gather benchmarks from `{}`: {err:?}",
-                                path.display()
-                            )
-                        })?;
-                        log::info!("Compiled {}", path.display());
-
-                        group = Some(BenchmarkGroup {
-                            binary: path,
-                            name: group_name.to_string(),
-                            benchmark_names: benchmarks,
-                        });
-                    }
+    for artifact in &mut cargo_iter {
+        let artifact = artifact?;
+        if let Some(ref executable) = artifact.executable {
+            // Found a binary compiled by a runtime benchmark crate.
+            // Execute it so that we find all the benchmarks it contains.
+            if artifact.target.kind.iter().any(|k| k == "bin") {
+                if group.is_some() {
+                    return Err(anyhow::anyhow!(
+                        "Runtime benchmark group `{group_name}` has produced multiple binaries"
+                    ));
                 }
+
+                let path = executable.as_std_path().to_path_buf();
+                let benchmarks = gather_benchmarks(&path).map_err(|err| {
+                    anyhow::anyhow!(
+                        "Cannot gather benchmarks from `{}`: {err:?}",
+                        path.display()
+                    )
+                })?;
+                log::info!("Compiled {}", path.display());
+
+                group = Some(BenchmarkGroup {
+                    binary: path,
+                    name: group_name.to_string(),
+                    benchmark_names: benchmarks,
+                });
             }
-            Message::TextLine(line) => {
-                println!("{line}")
-            }
-            Message::CompilerMessage(msg) => {
-                let message = msg.message.rendered.unwrap_or(msg.message.message);
-                messages.push_str(&message);
-                print!("{message}");
-            }
-            _ => {}
         }
     }
-
-    let output = cargo_process.wait()?;
-    if !output.success() {
-        Err(anyhow::anyhow!(
-            "Failed to compile runtime benchmark, exit code {}\n{messages}",
-            output.code().unwrap_or(1),
-        ))
-    } else {
-        let group = group.ok_or_else(|| {
-            anyhow::anyhow!("Runtime benchmark group `{group_name}` has not produced any binary")
-        })?;
-        Ok(group)
-    }
+    cargo_iter
+        .finish()
+        .with_context(|| format!("Failed to compile runtime benchmark `{group_name}`"))?;
+    let group = group.ok_or_else(|| {
+        anyhow::anyhow!("Runtime benchmark group `{group_name}` has not produced any binary")
+    })?;
+    Ok(group)
 }
 
 /// Starts the compilation of a single runtime benchmark crate.
@@ -345,18 +328,13 @@ fn start_cargo_build(
     benchmark_dir: &Path,
     target_dir: Option<&Path>,
     opts: &RuntimeCompilationOpts,
-) -> anyhow::Result<Child> {
+) -> anyhow::Result<CargoArtifactIter> {
     let mut command = Command::new(&toolchain.components.cargo);
     command
         .env("RUSTC", &toolchain.components.rustc)
         .arg("build")
         .arg("--release")
-        .arg("--message-format")
-        .arg("json-diagnostic-short")
-        .current_dir(benchmark_dir)
-        .stdin(Stdio::null())
-        .stdout(Stdio::piped())
-        .stderr(Stdio::null());
+        .current_dir(benchmark_dir);
 
     if let Some(ref debug_info) = opts.debug_info {
         command.env("CARGO_PROFILE_RELEASE_DEBUG", debug_info);
@@ -371,10 +349,8 @@ fn start_cargo_build(
     #[cfg(feature = "precise-cachegrind")]
     command.arg("--features").arg("benchlib/precise-cachegrind");
 
-    let child = command
-        .spawn()
-        .map_err(|error| anyhow::anyhow!("Failed to start cargo: {:?}", error))?;
-    Ok(child)
+    CargoArtifactIter::from_cargo_cmd(command)
+        .map_err(|error| anyhow::anyhow!("Failed to start cargo: {:?}", error))
 }
 
 /// Uses a command from `benchlib` to find the benchmark names from the given