build-manifest: refactor detecting package versions

4 files changed, 118 insertions, 132 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 5a2dd23a7ae..f18db340aa6 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -231,8 +231,10 @@ name = "build-manifest"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "flate2",
  "serde",
  "serde_json",
+ "tar",
  "toml",
 ]
 
diff --git a/src/tools/build-manifest/Cargo.toml b/src/tools/build-manifest/Cargo.toml
index 2da224a54dc..4f89c31936d 100644
--- a/src/tools/build-manifest/Cargo.toml
+++ b/src/tools/build-manifest/Cargo.toml
@@ -9,3 +9,5 @@ toml = "0.5"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 anyhow = "1.0.32"
+flate2 = "1.0.16"
+tar = "0.4.29"
diff --git a/src/tools/build-manifest/src/main.rs b/src/tools/build-manifest/src/main.rs
index c694948bac0..d8813f95be4 100644
--- a/src/tools/build-manifest/src/main.rs
+++ b/src/tools/build-manifest/src/main.rs
@@ -236,24 +236,6 @@ struct Builder {
     s3_address: String,
     date: String,
 
-    rust_version: Option<String>,
-    cargo_version: Option<String>,
-    rls_version: Option<String>,
-    rust_analyzer_version: Option<String>,
-    clippy_version: Option<String>,
-    rustfmt_version: Option<String>,
-    llvm_tools_version: Option<String>,
-    miri_version: Option<String>,
-
-    rust_git_commit_hash: Option<String>,
-    cargo_git_commit_hash: Option<String>,
-    rls_git_commit_hash: Option<String>,
-    rust_analyzer_git_commit_hash: Option<String>,
-    clippy_git_commit_hash: Option<String>,
-    rustfmt_git_commit_hash: Option<String>,
-    llvm_tools_git_commit_hash: Option<String>,
-    miri_git_commit_hash: Option<String>,
-
     should_sign: bool,
 }
 
@@ -286,7 +268,7 @@ fn main() {
     }
 
     Builder {
-        versions: Versions::new(&channel, Path::new(&monorepo_path)).unwrap(),
+        versions: Versions::new(&channel, &input, Path::new(&monorepo_path)).unwrap(),
 
         input,
         output,
@@ -295,24 +277,6 @@ fn main() {
         s3_address,
         date,
 
-        rust_version: None,
-        cargo_version: None,
-        rls_version: None,
-        rust_analyzer_version: None,
-        clippy_version: None,
-        rustfmt_version: None,
-        llvm_tools_version: None,
-        miri_version: None,
-
-        rust_git_commit_hash: None,
-        cargo_git_commit_hash: None,
-        rls_git_commit_hash: None,
-        rust_analyzer_git_commit_hash: None,
-        clippy_git_commit_hash: None,
-        rustfmt_git_commit_hash: None,
-        llvm_tools_git_commit_hash: None,
-        miri_git_commit_hash: None,
-
         should_sign,
     }
     .build();
@@ -320,26 +284,6 @@ fn main() {
 
 impl Builder {
     fn build(&mut self) {
-        self.rust_version = self.version("rust", "x86_64-unknown-linux-gnu");
-        self.cargo_version = self.version("cargo", "x86_64-unknown-linux-gnu");
-        self.rls_version = self.version("rls", "x86_64-unknown-linux-gnu");
-        self.rust_analyzer_version = self.version("rust-analyzer", "x86_64-unknown-linux-gnu");
-        self.clippy_version = self.version("clippy", "x86_64-unknown-linux-gnu");
-        self.rustfmt_version = self.version("rustfmt", "x86_64-unknown-linux-gnu");
-        self.llvm_tools_version = self.version("llvm-tools", "x86_64-unknown-linux-gnu");
-        self.miri_version = self.version("miri", "x86_64-unknown-linux-gnu");
-
-        self.rust_git_commit_hash = self.git_commit_hash("rust", "x86_64-unknown-linux-gnu");
-        self.cargo_git_commit_hash = self.git_commit_hash("cargo", "x86_64-unknown-linux-gnu");
-        self.rls_git_commit_hash = self.git_commit_hash("rls", "x86_64-unknown-linux-gnu");
-        self.rust_analyzer_git_commit_hash =
-            self.git_commit_hash("rust-analyzer", "x86_64-unknown-linux-gnu");
-        self.clippy_git_commit_hash = self.git_commit_hash("clippy", "x86_64-unknown-linux-gnu");
-        self.rustfmt_git_commit_hash = self.git_commit_hash("rustfmt", "x86_64-unknown-linux-gnu");
-        self.llvm_tools_git_commit_hash =
-            self.git_commit_hash("llvm-tools", "x86_64-unknown-linux-gnu");
-        self.miri_git_commit_hash = self.git_commit_hash("miri", "x86_64-unknown-linux-gnu");
-
         self.check_toolstate();
         self.digest_and_sign();
         let manifest = self.build_manifest();
@@ -368,8 +312,7 @@ impl Builder {
         // Mark some tools as missing based on toolstate.
         if toolstates.get("miri").map(|s| &*s as &str) != Some("test-pass") {
             println!("Miri tests are not passing, removing component");
-            self.miri_version = None;
-            self.miri_git_commit_hash = None;
+            self.versions.disable_version(&PkgType::Miri);
         }
     }
 
@@ -471,13 +414,10 @@ impl Builder {
     }
 
     fn rust_package(&mut self, manifest: &Manifest) -> Package {
+        let version_info = self.versions.version(&PkgType::Rust).expect("missing Rust tarball");
         let mut pkg = Package {
-            version: self
-                .cached_version("rust")
-                .as_ref()
-                .expect("Couldn't find Rust version")
-                .clone(),
-            git_commit_hash: self.cached_git_commit_hash("rust").clone(),
+            version: version_info.version.expect("missing Rust version"),
+            git_commit_hash: version_info.git_commit,
             target: BTreeMap::new(),
         };
         for host in HOSTS {
@@ -583,12 +523,11 @@ impl Builder {
     }
 
     fn package(&mut self, pkgname: &str, dst: &mut BTreeMap<String, Package>, targets: &[&str]) {
-        let (version, mut is_present) = self
-            .cached_version(pkgname)
-            .as_ref()
-            .cloned()
-            .map(|version| (version, true))
-            .unwrap_or_default(); // `is_present` defaults to `false` here.
+        let version_info = self
+            .versions
+            .version(&PkgType::from_component(pkgname))
+            .expect("failed to load package version");
+        let mut is_present = version_info.present;
 
         // Never ship nightly-only components for other trains.
         if self.versions.channel() != "nightly" && NIGHTLY_ONLY_COMPONENTS.contains(&pkgname) {
@@ -635,8 +574,8 @@ impl Builder {
         dst.insert(
             pkgname.to_string(),
             Package {
-                version,
-                git_commit_hash: self.cached_git_commit_hash(pkgname).clone(),
+                version: version_info.version.unwrap_or_default(),
+                git_commit_hash: version_info.git_commit,
                 target: targets,
             },
         );
@@ -646,64 +585,6 @@ impl Builder {
         format!("{}/{}/{}", self.s3_address, self.date, filename)
     }
 
-    fn cached_version(&self, component: &str) -> &Option<String> {
-        use PkgType::*;
-        match PkgType::from_component(component) {
-            Cargo => &self.cargo_version,
-            Rls => &self.rls_version,
-            RustAnalyzer => &self.rust_analyzer_version,
-            Clippy => &self.clippy_version,
-            Rustfmt => &self.rustfmt_version,
-            LlvmTools => &self.llvm_tools_version,
-            Miri => &self.miri_version,
-            _ => &self.rust_version,
-        }
-    }
-
-    fn cached_git_commit_hash(&self, component: &str) -> &Option<String> {
-        use PkgType::*;
-        match PkgType::from_component(component) {
-            Cargo => &self.cargo_git_commit_hash,
-            Rls => &self.rls_git_commit_hash,
-            RustAnalyzer => &self.rust_analyzer_git_commit_hash,
-            Clippy => &self.clippy_git_commit_hash,
-            Rustfmt => &self.rustfmt_git_commit_hash,
-            LlvmTools => &self.llvm_tools_git_commit_hash,
-            Miri => &self.miri_git_commit_hash,
-            _ => &self.rust_git_commit_hash,
-        }
-    }
-
-    fn version(&mut self, component: &str, target: &str) -> Option<String> {
-        self.untar(component, target, |filename| format!("{}/version", filename))
-    }
-
-    fn git_commit_hash(&mut self, component: &str, target: &str) -> Option<String> {
-        self.untar(component, target, |filename| format!("{}/git-commit-hash", filename))
-    }
-
-    fn untar<F>(&mut self, component: &str, target: &str, dir: F) -> Option<String>
-    where
-        F: FnOnce(String) -> String,
-    {
-        let filename = self
-            .versions
-            .tarball_name(&PkgType::from_component(component), target)
-            .expect("failed to retrieve the tarball path");
-
-        let mut cmd = Command::new("tar");
-        cmd.arg("xf")
-            .arg(self.input.join(&filename))
-            .arg(dir(filename.replace(".tar.gz", "")))
-            .arg("-O");
-        let output = t!(cmd.output());
-        if output.status.success() {
-            Some(String::from_utf8_lossy(&output.stdout).trim().to_string())
-        } else {
-            None
-        }
-    }
-
     fn hash(&self, path: &Path) -> String {
         let sha = t!(Command::new("shasum")
             .arg("-a")
diff --git a/src/tools/build-manifest/src/versions.rs b/src/tools/build-manifest/src/versions.rs
index 151cfa0d9c6..3b5caf1cd7b 100644
--- a/src/tools/build-manifest/src/versions.rs
+++ b/src/tools/build-manifest/src/versions.rs
@@ -1,6 +1,12 @@
 use anyhow::{Context, Error};
+use flate2::read::GzDecoder;
 use std::collections::HashMap;
+use std::fs::File;
+use std::io::Read;
 use std::path::{Path, PathBuf};
+use tar::Archive;
+
+const DEFAULT_TARGET: &str = "x86_64-unknown-linux-gnu";
 
 #[derive(Debug, Hash, Eq, PartialEq, Clone)]
 pub(crate) enum PkgType {
@@ -61,17 +67,46 @@ impl PkgType {
             PkgType::Other(component) => component,
         }
     }
+
+    fn should_use_rust_version(&self) -> bool {
+        match self {
+            PkgType::Cargo => false,
+            PkgType::Rls => false,
+            PkgType::RustAnalyzer => false,
+            PkgType::Clippy => false,
+            PkgType::Rustfmt => false,
+            PkgType::LlvmTools => false,
+            PkgType::Miri => false,
+
+            PkgType::Rust => true,
+            PkgType::RustSrc => true,
+            PkgType::Other(_) => true,
+        }
+    }
+}
+
+#[derive(Debug, Default, Clone)]
+pub(crate) struct VersionInfo {
+    pub(crate) version: Option<String>,
+    pub(crate) git_commit: Option<String>,
+    pub(crate) present: bool,
 }
 
 pub(crate) struct Versions {
     channel: String,
     rustc_version: String,
     monorepo_root: PathBuf,
+    dist_path: PathBuf,
     package_versions: HashMap<PkgType, String>,
+    versions: HashMap<PkgType, VersionInfo>,
 }
 
 impl Versions {
-    pub(crate) fn new(channel: &str, monorepo_root: &Path) -> Result<Self, Error> {
+    pub(crate) fn new(
+        channel: &str,
+        dist_path: &Path,
+        monorepo_root: &Path,
+    ) -> Result<Self, Error> {
         Ok(Self {
             channel: channel.into(),
             rustc_version: std::fs::read_to_string(monorepo_root.join("src").join("version"))
@@ -79,7 +114,9 @@ impl Versions {
                 .trim()
                 .to_string(),
             monorepo_root: monorepo_root.into(),
+            dist_path: dist_path.into(),
             package_versions: HashMap::new(),
+            versions: HashMap::new(),
         })
     }
 
@@ -87,6 +124,70 @@ impl Versions {
         &self.channel
     }
 
+    pub(crate) fn version(&mut self, mut package: &PkgType) -> Result<VersionInfo, Error> {
+        if package.should_use_rust_version() {
+            package = &PkgType::Rust;
+        }
+
+        match self.versions.get(package) {
+            Some(version) => Ok(version.clone()),
+            None => {
+                let version_info = self.load_version_from_tarball(package)?;
+                self.versions.insert(package.clone(), version_info.clone());
+                Ok(version_info)
+            }
+        }
+    }
+
+    fn load_version_from_tarball(&mut self, package: &PkgType) -> Result<VersionInfo, Error> {
+        let tarball_name = self.tarball_name(package, DEFAULT_TARGET)?;
+        let tarball = self.dist_path.join(tarball_name);
+
+        let file = match File::open(&tarball) {
+            Ok(file) => file,
+            Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
+                // Missing tarballs do not return an error, but return empty data.
+                return Ok(VersionInfo::default());
+            }
+            Err(err) => return Err(err.into()),
+        };
+        let mut tar = Archive::new(GzDecoder::new(file));
+
+        let mut version = None;
+        let mut git_commit = None;
+        for entry in tar.entries()? {
+            let mut entry = entry?;
+
+            let dest;
+            match entry.path()?.components().nth(1).and_then(|c| c.as_os_str().to_str()) {
+                Some("version") => dest = &mut version,
+                Some("git-commit-hash") => dest = &mut git_commit,
+                _ => continue,
+            }
+            let mut buf = String::new();
+            entry.read_to_string(&mut buf)?;
+            *dest = Some(buf);
+
+            // Short circuit to avoid reading the whole tar file if not necessary.
+            if version.is_some() && git_commit.is_some() {
+                break;
+            }
+        }
+
+        Ok(VersionInfo { version, git_commit, present: true })
+    }
+
+    pub(crate) fn disable_version(&mut self, package: &PkgType) {
+        match self.versions.get_mut(package) {
+            Some(version) => {
+                *version = VersionInfo::default();
+            }
+            None => {
+                self.versions.insert(package.clone(), VersionInfo::default());
+            }
+        }
+    }
+
     pub(crate) fn tarball_name(
         &mut self,
         package: &PkgType,