Add command for downloading a new benchmark

Author: Kobzol

Cargo.lock

@@ -31,6 +31,7 @@ crossbeam-utils = "0.8"
 snap = "1"
 filetime = "0.2.14"
 walkdir = "2"
+flate2 = { version = "1.0.22", features = ["rust_backend"] }
 
 [target.'cfg(windows)'.dependencies]
 miow = "0.3"

collector/Cargo.toml

@@ -25,7 +25,7 @@ use tokio::runtime::Runtime;
 mod rustc;
 
 #[cfg(windows)]
-fn rename<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q) -> anyhow::Result<()> {
+pub fn rename<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q) -> anyhow::Result<()> {
     let (from, to) = (from.as_ref(), to.as_ref());
 
     let ctx = format!("renaming file {:?} to {:?}", from, to);
@@ -38,7 +38,7 @@ fn rename<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q) -> anyhow::Result<()>
 }
 
 #[cfg(unix)]
-fn rename<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q) -> anyhow::Result<()> {
+pub fn rename<P: AsRef<Path>, Q: AsRef<Path>>(from: P, to: Q) -> anyhow::Result<()> {
     let (from, to) = (from.as_ref(), to.as_ref());
     if fs::rename(from, to).is_err() {
         // This is necessary if from and to are on different

collector/src/execute.rs

@@ -6,6 +6,7 @@ use database::{ArtifactId, Commit};
 use log::debug;
 use std::collections::HashSet;
 use std::fs;
+use std::fs::File;
 use std::path::{Path, PathBuf};
 use std::process;
 use std::process::{Command, Stdio};
@@ -16,6 +17,7 @@ mod execute;
 mod sysroot;
 
 use execute::{BenchProcessor, Benchmark, BenchmarkName, ProfileProcessor, Profiler};
+use std::io::Write;
 use sysroot::Sysroot;
 
 #[derive(Debug, Copy, Clone)]
@@ -848,6 +850,27 @@ enum Commands {
 
     /// Installs the next commit for perf.rust-lang.org
     InstallNext,
+
+    /// Download a crate into collector/benchmarks.
+    Download(DownloadCommand),
+}
+
+#[derive(Debug, clap::Parser)]
+struct DownloadCommand {
+    /// Name of the benchmark created directory
+    #[clap(long, global = true)]
+    name: Option<String>,
+
+    #[clap(subcommand)]
+    command: DownloadSubcommand,
+}
+
+#[derive(Debug, clap::Parser)]
+enum DownloadSubcommand {
+    /// Download a crate from a git repository.
+    Git { url: String },
+    /// Download a crate from crates.io.
+    Crate { krate: String, version: String },
 }
 
 fn main_result() -> anyhow::Result<i32> {
@@ -1141,7 +1164,99 @@ fn main_result() -> anyhow::Result<i32> {
 
             Ok(0)
         }
+        Commands::Download(opts) => {
+            let crate_path = match opts.command {
+                DownloadSubcommand::Git { url } => {
+                    download_from_git(&benchmark_dir, opts.name.as_deref(), &url)?
+                }
+                DownloadSubcommand::Crate { krate, version } => {
+                    download_from_crates_io(&benchmark_dir, &krate, &version, opts.name.as_deref())?
+                }
+            };
+            prepare_downloaded_benchmark(&crate_path);
+            Ok(0)
+        }
+    }
+}
+
+fn download_from_git(
+    benchmark_dir: &Path,
+    name: Option<&str>,
+    url: &str,
+) -> anyhow::Result<PathBuf> {
+    let target = PathBuf::from(benchmark_dir);
+    let crate_name = match name {
+        Some(name) => name,
+        None => url
+            .trim_end_matches("/")
+            .trim_end_matches(".git")
+            .split("/")
+            .last()
+            .expect("Crate name could not be determined from git URL"),
+    };
+    let target = target.join(crate_name);
+    if target.exists() {
+        return Err(anyhow::anyhow!(
+            "Directory {} already exists",
+            target.display()
+        ));
+    }
+
+    Command::new("git")
+        .arg("clone")
+        .arg(url)
+        .arg(&target)
+        .output()
+        .expect("Git clone failed");
+    Ok(target)
+}
+
+fn download_from_crates_io(
+    benchmark_dir: &Path,
+    krate: &str,
+    version: &str,
+    name: Option<&str>,
+) -> anyhow::Result<PathBuf> {
+    let full_name = format!("{krate}-{version}");
+    let name = name.unwrap_or(&full_name);
+    let target = PathBuf::from(benchmark_dir).join(name);
+    if target.exists() {
+        return Err(anyhow::anyhow!(
+            "Directory {} already exists",
+            target.display()
+        ));
+    }
+
+    let url = format!("https://crates.io/api/v1/crates/{krate}/{version}/download");
+    let body = reqwest::blocking::get(url).expect("Cannot download crate");
+    let data = flate2::read::GzDecoder::new(body);
+    let mut archive = tar::Archive::new(data);
+
+    let tmpdir = tempfile::TempDir::new().unwrap();
+    archive.unpack(&tmpdir)?;
+    let unpacked_dir = std::fs::read_dir(&tmpdir.path())
+        .unwrap()
+        .next()
+        .unwrap()
+        .unwrap()
+        .path();
+    execute::rename(&unpacked_dir, &target)?;
+
+    Ok(target)
+}
+
+fn prepare_downloaded_benchmark(path: &Path) {
+    let manifest_path = path.join("Cargo.toml");
+    let manifest = std::fs::read_to_string(&manifest_path).expect("Cannot open Cargo.toml");
+    if !manifest.contains("[workspace]") {
+        let mut file = File::options().append(true).open(&manifest_path).unwrap();
+        writeln!(file, "[workspace]").expect("Cannot append workspace to Cargo.toml");
     }
+    Command::new("cargo")
+        .arg("generate-lockfile")
+        .current_dir(manifest_path.parent().unwrap())
+        .status()
+        .expect("Cannot generate lock file");
 }
 
 pub fn get_commit_or_fake_it(sha: &str) -> anyhow::Result<Commit> {