benches/capture/src/main.rs

   1 //! This tool helps to capture the `Cargo.toml` files of a workspace.
   2 //!
   3 //! Run it by passing a list of workspaces to capture.
   4 //! Use the `-f` flag to allow it to overwrite existing captures.
   5 //! The workspace will be saved in a `.tgz` file in the `../workspaces` directory.
   6
   7 use flate2::{Compression, GzBuilder};
   8 use std::fs;
   9 use std::path::{Path, PathBuf};
  10 use std::process::Command;
  11
  12 fn main() {
  13     let force = std::env::args().any(|arg| arg == "-f");
  14     let dest = Path::new(env!("CARGO_MANIFEST_DIR"))
  15         .parent()
  16         .unwrap()
  17         .join("workspaces");
  18     if !dest.exists() {
  19         panic!("expected {} to exist", dest.display());
  20     }
  21     for arg in std::env::args().skip(1).filter(|arg| !arg.starts_with("-")) {
  22         let source_root = fs::canonicalize(arg).unwrap();
  23         capture(&source_root, &dest, force);
  24     }
  25 }
  26
  27 fn capture(source_root: &Path, dest: &Path, force: bool) {
  28     let name = Path::new(source_root.file_name().unwrap());
  29     let mut dest_gz = PathBuf::from(dest);
  30     dest_gz.push(name);
  31     dest_gz.set_extension("tgz");
  32     if dest_gz.exists() {
  33         if !force {
  34             panic!(
  35                 "dest {:?} already exists, use -f to force overwriting",
  36                 dest_gz
  37             );
  38         }
  39         fs::remove_file(&dest_gz).unwrap();
  40     }
  41     let vcs_info = capture_vcs_info(source_root, force);
  42     let dst = fs::File::create(&dest_gz).unwrap();
  43     let encoder = GzBuilder::new()
  44         .filename(format!("{}.tar", name.to_str().unwrap()))
  45         .write(dst, Compression::best());
  46     let mut ar = tar::Builder::new(encoder);
  47     ar.mode(tar::HeaderMode::Deterministic);
  48     if let Some(info) = &vcs_info {
  49         add_ar_file(&mut ar, &name.join(".cargo_vcs_info.json"), info);
  50     }
  51
  52     // Gather all local packages.
  53     let metadata = cargo_metadata::MetadataCommand::new()
  54         .manifest_path(source_root.join("Cargo.toml"))
  55         .features(cargo_metadata::CargoOpt::AllFeatures)
  56         .exec()
  57         .expect("cargo_metadata failed");
  58     let mut found_root = false;
  59     for package in &metadata.packages {
  60         if package.source.is_some() {
  61             continue;
  62         }
  63         let manifest_path = package.manifest_path.as_std_path();
  64         copy_manifest(&manifest_path, &mut ar, name, &source_root);
  65         found_root |= manifest_path == source_root.join("Cargo.toml");
  66     }
  67     if !found_root {
  68         // A virtual workspace.
  69         let contents = fs::read_to_string(source_root.join("Cargo.toml")).unwrap();
  70         assert!(!contents.contains("[package]"));
  71         add_ar_file(&mut ar, &name.join("Cargo.toml"), &contents);
  72     }
  73     let lock = fs::read_to_string(source_root.join("Cargo.lock")).unwrap();
  74     add_ar_file(&mut ar, &name.join("Cargo.lock"), &lock);
  75     let encoder = ar.into_inner().unwrap();
  76     encoder.finish().unwrap();
  77     eprintln!("created {}", dest_gz.display());
  78 }
  79
  80 fn copy_manifest<W: std::io::Write>(
  81     manifest_path: &Path,
  82     ar: &mut tar::Builder<W>,
  83     name: &Path,
  84     source_root: &Path,
  85 ) {
  86     let relative_path = manifest_path
  87         .parent()
  88         .unwrap()
  89         .strip_prefix(source_root)
  90         .expect("workspace member should be under workspace root");
  91     let relative_path = name.join(relative_path);
  92     let contents = fs::read_to_string(&manifest_path).unwrap();
  93     let mut manifest: toml::Value = toml::from_str(&contents).unwrap();
  94     let remove = |obj: &mut toml::Value, name| {
  95         let table = obj.as_table_mut().unwrap();
  96         if table.contains_key(name) {
  97             table.remove(name);
  98         }
  99     };
 100     remove(&mut manifest, "lib");
 101     remove(&mut manifest, "bin");
 102     remove(&mut manifest, "example");
 103     remove(&mut manifest, "test");
 104     remove(&mut manifest, "bench");
 105     remove(&mut manifest, "profile");
 106     if let Some(package) = manifest.get_mut("package") {
 107         remove(package, "default-run");
 108     }
 109     let contents = toml::to_string(&manifest).unwrap();
 110     add_ar_file(ar, &relative_path.join("Cargo.toml"), &contents);
 111     add_ar_file(ar, &relative_path.join("src").join("lib.rs"), "");
 112 }
 113
 114 fn add_ar_file<W: std::io::Write>(ar: &mut tar::Builder<W>, path: &Path, contents: &str) {
 115     let mut header = tar::Header::new_gnu();
 116     header.set_entry_type(tar::EntryType::file());
 117     header.set_mode(0o644);
 118     header.set_size(contents.len() as u64);
 119     header.set_mtime(123456789);
 120     header.set_cksum();
 121     ar.append_data(&mut header, path, contents.as_bytes())
 122         .unwrap();
 123 }
 124
 125 fn capture_vcs_info(ws_root: &Path, force: bool) -> Option<String> {
 126     let maybe_git = |command: &str| {
 127         Command::new("git")
 128             .current_dir(ws_root)
 129             .args(command.split_whitespace().collect::<Vec<_>>())
 130             .output()
 131             .expect("git should be installed")
 132     };
 133     assert!(ws_root.join("Cargo.toml").exists());
 134     let relative = maybe_git("ls-files --full-name Cargo.toml");
 135     if !relative.status.success() {
 136         if !force {
 137             panic!("git repository not detected, use -f to force");
 138         }
 139         return None;
 140     }
 141     let p = Path::new(std::str::from_utf8(&relative.stdout).unwrap().trim());
 142     let relative = p.parent().unwrap();
 143     if !force {
 144         let has_changes = !maybe_git("diff-index --quiet HEAD .").status.success();
 145         if has_changes {
 146             panic!("git repo appears to have changes, use -f to force, or clean the repo");
 147         }
 148     }
 149     let commit = maybe_git("rev-parse HEAD");
 150     assert!(commit.status.success());
 151     let commit = std::str::from_utf8(&commit.stdout).unwrap().trim();
 152     let remote = maybe_git("remote get-url origin");
 153     assert!(remote.status.success());
 154     let remote = std::str::from_utf8(&remote.stdout).unwrap().trim();
 155     let info = format!(
 156         "{{\n  \"git\": {{\n    \"sha1\": \"{}\",\n     \"remote\": \"{}\"\n  }},\
 157          \n  \"path_in_vcs\": \"{}\"\n}}\n",
 158         commit,
 159         remote,
 160         relative.display()
 161     );
 162     eprintln!("recording vcs info:\n{}", info);
 163     Some(info)
 164 }