Add binary cachingbinary-caching

11 files changed, 491 insertions, 16 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 8e1126df6..2344293b3 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1,5 +1,7 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
+version = 3
+
 [[package]]
 name = "ahash"
 version = "0.4.7"
@@ -395,6 +397,15 @@ dependencies = [
 ]
 
 [[package]]
+name = "dirs"
+version = "3.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "30baa043103c9d0c2a57cf537cc2f35623889dc0d405e6c3cccfadbc81c71309"
+dependencies = [
+ "dirs-sys",
+]
+
+[[package]]
 name = "dirs-next"
 version = "2.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -405,6 +416,17 @@ dependencies = [
 ]
 
 [[package]]
+name = "dirs-sys"
+version = "0.3.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "03d86534ed367a67548dc68113a0f5db55432fdfbb6e6f9d77704397d95d5780"
+dependencies = [
+ "libc",
+ "redox_users 0.4.0",
+ "winapi",
+]
+
+[[package]]
 name = "dirs-sys-next"
 version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -861,7 +883,7 @@ checksum = "3dfb6b71a9a89cd38b395d994214297447e8e63b1ba5708a9a2b0b1048ceda76"
 dependencies = [
  "cc",
  "chrono",
- "dirs",
+ "dirs 1.0.5",
  "objc-foundation",
 ]
 
@@ -1718,6 +1740,7 @@ dependencies = [
  "serde_json",
  "shadow-rs",
  "shell-words",
+ "starship-cache",
  "starship_module_config_derive",
  "strsim 0.10.0",
  "sys-info",
@@ -1734,6 +1757,17 @@ dependencies = [
 ]
 
 [[package]]
+name = "starship-cache"
+version = "0.1.0"
+dependencies = [
+ "dirs 3.0.2",
+ "serde",
+ "tempfile",
+ "thiserror",
+ "toml",
+]
+
+[[package]]
 name = "starship_module_config_derive"
 version = "0.2.1"
 dependencies = [
@@ -1857,6 +1891,26 @@ dependencies = [
 ]
 
 [[package]]
+name = "thiserror"
+version = "1.0.25"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fa6f76457f59514c7eeb4e59d891395fab0b2fd1d40723ae737d64153392e9c6"
+dependencies = [
+ "thiserror-impl",
+]
+
+[[package]]
+name = "thiserror-impl"
+version = "1.0.25"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8a36768c0fbf1bb15eca10defa29526bda730a2376c2ab4393ccfa16fb1a318d"
+dependencies = [
+ "proc-macro2",
+ "quote 1.0.9",
+ "syn 1.0.72",
+]
+
+[[package]]
 name = "time"
 version = "0.1.44"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index 1a2ad5e66..67f911ca3 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -49,7 +49,8 @@ once_cell = "1.8.0"
 chrono = "0.4.19"
 sys-info = "0.9.0"
 byte-unit = "4.0.12"
-starship_module_config_derive = { version = "0.2.1", path = "starship_module_config_derive" }
+starship_module_config_derive = { version = "0.2.1", path = "crates/starship_module_config_derive" }
+starship-cache = { version = "0.1.0", path = "crates/starship-cache" }
 yaml-rust = "0.4.5"
 pest = "2.1.3"
 pest_derive = "2.1.0"
@@ -103,3 +104,6 @@ lto = true
 [[bin]]
 name = "starship"
 path = "src/main.rs"
+
+[workspace]
+members = ["crates/*"]
diff --git a/crates/starship-cache/Cargo.toml b/crates/starship-cache/Cargo.toml
new file mode 100644
index 000000000..bb66cd150
--- /dev/null
+++ b/crates/starship-cache/Cargo.toml
@@ -0,0 +1,18 @@
+[package]
+name = "starship-cache"
+version = "0.1.0"
+authors = ["Starship Contributors"]
+description = "Intelligent caching for Starship"
+edition = "2018"
+license = "ISC"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+dirs = "3.0.2"
+serde = { version = "1.0.126", features = ["derive"] }
+thiserror = "1.0.25"
+toml = "0.5.8"
+
+[dev-dependencies]
+tempfile = "3.2.0"
diff --git a/crates/starship-cache/src/errors.rs b/crates/starship-cache/src/errors.rs
new file mode 100644
index 000000000..248a3af3b
--- /dev/null
+++ b/crates/starship-cache/src/errors.rs
@@ -0,0 +1,17 @@
+use std::io;
+
+#[non_exhaustive]
+#[derive(thiserror::Error, Debug)]
+pub enum Error {
+    #[error("cannot open cache file")]
+    OpenFile(#[source] io::Error),
+
+    #[error("cannot write cache file")]
+    WriteFile(#[source] io::Error),
+
+    #[error("cannot read binary metadata")]
+    ReadMetadata(#[source] io::Error),
+
+    #[error("unable to serialize cache")]
+    SerializeCache(#[source] toml::ser::Error),
+}
diff --git a/crates/starship-cache/src/lib.rs b/crates/starship-cache/src/lib.rs
new file mode 100644
index 000000000..77fa0bb76
--- /dev/null
+++ b/crates/starship-cache/src/lib.rs
@@ -0,0 +1,312 @@
+//! The on-disk caching functionality for Starship.
+//!
+//! This module contains the caching mechanism allowing Starship to reuse the
+//! output of previously run commands when possible.
+//!
+//! The cache stores the output of commands, and the metadata of the binaries
+//! being called at the time the command is run. When the binary's metadata
+//! changes, the cache clears all the values of the commands calling that binary.
+//!
+//! The goals of this library are to be quick to cache outputs, quick to retreive
+//! cached values, compatible with version-managed tools, and easy to troubleshoot.
+
+pub mod errors;
+
+pub use errors::Error;
+use serde::{Deserialize, Serialize};
+use std::{
+    collections::HashMap,
+    convert::TryFrom,
+    fs::{self, OpenOptions},
+    io::Read,
+    path::{Path, PathBuf},
+    process::Output,
+    time::UNIX_EPOCH,
+};
+
+type FullCommand = String;
+type BinaryPath = PathBuf;
+
+const CURRENT_VERSION: u8 = 1;
+
+/// An instance of the binary output cache
+pub struct Cache {
+    /// The path of the cache file the cache serializes to
+    path: PathBuf,
+    /// Whether the cache has been changed and requires writing to disk
+    changed: bool,
+    /// The cache's internal state
+    contents: CacheContents,
+}
+
+impl Cache {
+    /// Create or parse a cache file at the given path
+    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
+        let mut file = OpenOptions::new()
+            .read(true)
+            .write(true)
+            .create(true)
+            .open(&path)
+            .map_err(Error::OpenFile)?;
+        let mut contents = String::new();
+
+        // Clear the cache if it is not valid UTF-8
+        file.read_to_string(&mut contents).unwrap_or_default();
+
+        // Clear the cache if it unable to be parsed
+        let mut cache: CacheContents = toml::from_str(&contents).unwrap_or_default();
+
+        // Clear the cache if it is not using the current version
+        if cache.version != CURRENT_VERSION {
+            cache = CacheContents::default();
+        }
+
+        Ok(Self {
+            path: path.as_ref().to_owned(),
+            changed: false,
+            contents: cache,
+        })
+    }
+
+    /// Get the output of the given command if it has been previously cached
+    pub fn get(&mut self, binary_path: &Path, command: &str) -> Option<&CachedOutput> {
+        let bin = self.contents.binaries.get(binary_path)?;
+
+        let current_metadata = BinaryMetadata::try_from(binary_path).ok()?;
+        let is_stale = current_metadata != bin.metadata;
+        if is_stale {
+            return None;
+        };
+
+        bin.commands.get(command)
+    }
+
+    /// Set the cached output of the given command
+    pub fn set<O: Into<CachedOutput>>(&mut self, binary_path: &Path, command: &str, output: O) {
+        let current_metadata = match BinaryMetadata::try_from(binary_path) {
+            Ok(metadata) => metadata,
+            // Skip caching if unable to read binary metadata
+            Err(_e) => return,
+        };
+        let mut bin = self
+            .contents
+            .binaries
+            .entry(binary_path.to_path_buf())
+            .or_insert(BinaryCache {
+                metadata: current_metadata.clone(),
+                commands: HashMap::new(),
+            });
+
+        let is_stale = current_metadata != bin.metadata;
+        if is_stale {
+            bin.metadata = current_metadata;
+            bin.commands.clear();
+        };
+
+        bin.commands.insert(command.to_owned(), output.into());
+        self.changed = true;
+    }
+
+    /// Write any cache updates to disk
+    pub fn write(&self) -> Result<(), Error> {
+        if !self.changed {
+            return Ok(());
+        };
+
+        let contents = toml::to_string(&self.contents).map_err(Error::SerializeCache)?;
+        fs::write(&self.path, contents).map_err(Error::WriteFile)?;
+        Ok(())
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+struct CacheContents {
+    /// The version of the cache file
+    version: u8,
+    /// A mapping of binaries' paths and their caches
+    binaries: HashMap<BinaryPath, BinaryCache>,
+}
+
+impl Default for CacheContents {
+    fn default() -> Self {
+        Self {
+            version: CURRENT_VERSION,
+            binaries: HashMap::new(),
+        }
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+struct BinaryCache {
+    /// The metadata of the binary at the time it was last called
+    /// If the binary's metadata changes, its cached data is cleared
+    metadata: BinaryMetadata,
+    /// A mapping of commands and their cached outputs
+    commands: HashMap<FullCommand, CachedOutput>,
+}
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct CachedOutput {
+    pub stdout: Vec<u8>,
+    pub stderr: Vec<u8>,
+    pub status: Option<i32>,
+}
+
+impl CachedOutput {
+    pub fn success(&self) -> bool {
+        self.status == Some(0)
+    }
+}
+
+impl From<Output> for CachedOutput {
+    fn from(output: Output) -> Self {
+        Self {
+            stdout: output.stdout,
+            stderr: output.stderr,
+            status: output.status.code()
+        }
+    }
+}
+
+#[derive(Serialize, Deserialize, Debug, PartialEq, Clone)]
+struct BinaryMetadata {
+    size: u64,
+    is_dir: bool,
+    is_file: bool,
+    readonly: bool,
+    c_time: u64,
+    m_time: u64,
+}
+
+impl TryFrom<&Path> for BinaryMetadata {
+    type Error = crate::Error;
+
+    fn try_from(path: &Path) -> Result<Self, Error> {
+        let metadata = fs::metadata(path).map_err(Error::ReadMetadata)?;
+
+        // If ctime or mtime are not provided, store `0` in their place
+        let c_time = match metadata.created() {
+            Err(_e) => 0,
+            Ok(t) => t
+                .duration_since(UNIX_EPOCH)
+                .map(|t| t.as_secs())
+                .unwrap_or(0),
+        };
+
+        let m_time = match metadata.modified() {
+            Err(_e) => 0,
+            Ok(t) => t
+                .duration_since(UNIX_EPOCH)
+                .map(|t| t.as_secs())
+                .unwrap_or(0),
+        };
+
+        Ok(Self {
+            size: metadata.len(),
+            is_dir: metadata.is_dir(),
+            is_file: metadata.is_file(),
+            readonly: metadata.permissions().readonly(),
+            c_time,
+            m_time,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::{fs::File, io::Write};
+    use tempfile::tempdir;
+
+    type Result = std::result::Result<(), Box<dyn std::error::Error>>;
+
+    #[test]
+    fn empty_cache_file_is_created() -> Result {
+        let dir = tempdir()?;
+        let cache_path = Path::join(dir.path(), "bin-cache");
+        let cache = Cache::new(&cache_path)?;
+        cache.write()?;
+
+        assert!(Path::exists(&cache_path));
+        Ok(())
+    }
+
+    #[test]
+    fn retreive_from_populated_cache() -> Result {
+        let dir = tempdir()?;
+        let cache_path = dir.path().join("bin-cache");
+        let mut cache = Cache::new(&cache_path)?;
+
+        // Create "node" binary
+        let bin_path = dir.path().join("node");
+        File::create(&bin_path)?;
+
+        // Populate cache with "node" output
+        let expected = "v14.16.0";
+        cache.set(&bin_path, "node --version", &expected);
+        cache.write()?;
+
+        // Retreive cached output
+        let mut new_cache = Cache::new(&cache_path)?;
+        let actual = new_cache.get(&bin_path, "node --version").unwrap();
+
+        assert_eq!(expected, actual);
+        Ok(())
+    }
+
+    #[test]
+    fn overrites_stale_cache() -> Result {
+        let dir = tempdir()?;
+        let cache_path = dir.path().join("bin-cache");
+        let mut cache = Cache::new(&cache_path)?;
+
+        // Create "node" binary
+        let bin_path = dir.path().join("node");
+        File::create(&bin_path)?;
+
+        // Populate cache with "node" output
+        let expected = "v14.16.0";
+        cache.set(&bin_path, "node -v", &expected);
+        cache.set(&bin_path, "node --help", &expected);
+        cache.set(&bin_path, "node --version", &expected);
+        cache.write()?;
+
+        // Update "node" binary
+        File::create(&bin_path)?.write(b"updated")?;
+
+        // Retreive cached output
+        let mut new_cache = Cache::new(&cache_path)?;
+
+        // Set a cached value again
+        new_cache.set(&bin_path, "node -v", "v15.0.0");
+
+        // The other, previously cached values, should be cleared as stale
+        assert_eq!(new_cache.get(&bin_path, "node --version"), None);
+        assert_eq!(new_cache.get(&bin_path, "node --help"), None);
+        Ok(())
+    }
+
+    #[test]
+    fn doesnt_retreive_stale_cache() -> Result {
+        let dir = tempdir()?;
+        let cache_path = dir.path().join("bin-cache");
+        let mut cache = Cache::new(&cache_path)?;
+
+        // Create "node" binary
+        let bin_path = dir.path().join("node");
+        File::create(&bin_path)?;
+
+        // Populate cache with "node" output
+        cache.set(&bin_path, "node --version", "v14.16.0");
+        cache.write()?;
+
+        // Update "node" binary
+        File::create(&bin_path)?.write(b"updated")?;
+
+        let mut new_cache = Cache::new(&cache_path)?;
+        let actual = new_cache.get(&bin_path, "node --version");
+
+        assert_eq!(None, actual);
+        Ok(())
+    }
+}
diff --git a/starship_module_config_derive/Cargo.toml b/crates/starship_module_config_derive/Cargo.toml
index 928eef5f6..928eef5f6 100644
--- a/starship_module_config_derive/Cargo.toml
+++ b/crates/starship_module_config_derive/Cargo.toml
diff --git a/starship_module_config_derive/LICENSE b/crates/starship_module_config_derive/LICENSE
index 1ce6f8232..1ce6f8232 100644
--- a/starship_module_config_derive/LICENSE
+++ b/crates/starship_module_config_derive/LICENSE
diff --git a/starship_module_config_derive/README.md b/crates/starship_module_config_derive/README.md
index b1a0ca9ee..b1a0ca9ee 100644
--- a/starship_module_config_derive/README.md
+++ b/crates/starship_module_config_derive/README.md
diff --git a/starship_module_config_derive/src/lib.rs b/crates/starship_module_config_derive/src/lib.rs
index 40d4301b9..40d4301b9 100644
--- a/starship_module_config_derive/src/lib.rs
+++ b/crates/starship_module_config_derive/src/lib.rs
diff --git a/src/modules/rust.rs b/src/modules/rust.rs
index 2dcf73087..2fd255b62 100644
--- a/src/modules/rust.rs
+++ b/src/modules/rust.rs
@@ -1,6 +1,7 @@
 use std::fs;
 use std::path::Path;
 use std::process::{Command, Output};
+use std::time::Duration;
 
 use serde::Deserialize;
 
@@ -8,6 +9,8 @@ use super::{Context, Module, RootModuleConfig};
 
 use crate::configs::rust::RustConfig;
 use crate::formatter::{StringFormatter, VersionFormatter};
+use crate::utils::exec_cmd;
+use starship_cache::CachedOutput;
 
 /// Creates a module with the current Rust version
 pub fn module<'a>(context: &'a Context) -> Option<Module<'a>> {
@@ -188,15 +191,20 @@ fn find_rust_toolchain_file(context: &Context) -> Option<String> {
 }
 
 fn execute_rustup_run_rustc_version(toolchain: &str) -> RustupRunRustcVersionOutcome {
-    Command::new("rustup")
-        .args(&["run", toolchain, "rustc", "--version"])
-        .output()
-        .map(extract_toolchain_from_rustup_run_rustc_version)
-        .unwrap_or(RustupRunRustcVersionOutcome::RustupNotWorking)
+    exec_cmd(
+        "rustup",
+        &["run", toolchain, "rustc", "--version"],
+        Duration::from_millis(500),
+    )
+    .map(extract_toolchain_from_rustup_run_rustc_version)
+    .unwrap_or(RustupRunRustcVersionOutcome::RustupNotWorking)
 }
 
-fn extract_toolchain_from_rustup_run_rustc_version(output: Output) -> RustupRunRustcVersionOutcome {
-    if output.status.success() {
+fn extract_toolchain_from_rustup_run_rustc_version<O: Into<CachedOutput>>(
+    output: O,
+) -> RustupRunRustcVersionOutcome {
+    let output = output.into();
+    if output.success() {
         if let Ok(output) = String::from_utf8(output.stdout) {
             return RustupRunRustcVersionOutcome::RustcVersion(output);
         }
@@ -212,9 +220,9 @@ fn extract_toolchain_from_rustup_run_rustc_version(output: Output) -> RustupRunR
 }
 
 fn execute_rustc_version() -> Option<String> {
-    match Command::new("rustc").arg("--version").output() {
-        Ok(output) => Some(String::from_utf8(output.stdout).unwrap()),
-        Err(_) => None,
+    match exec_cmd("rustc", &["--version"], Duration::from_millis(500)) {
+        Some(output) => Some(output.stdout),
+        None => None,
     }
 }
 
diff --git a/src/utils.rs b/src/utils.rs
index 67265c9c8..22b7a952b 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -1,8 +1,10 @@
 use process_control::{ChildExt, Timeout};
+use starship_cache::{Cache, CachedOutput};
+use std::convert::TryInto;
 use std::fmt::Debug;
 use std::fs::read_to_string;
 use std::io::Result;
-use std::path::Path;
+use std::path::{Path, PathBuf};
 use std::process::{Command, Stdio};
 use std::time::{Duration, Instant};
 
@@ -35,6 +37,16 @@ impl PartialEq for CommandOutput {
     }
 }
 
+impl From<CommandOutput> for CachedOutput {
+    fn from(output: CommandOutput) -> Self {
+        Self {
+            stdout: output.stdout.into_bytes(),
+            stderr: output.stderr.into_bytes(),
+            status: Some(0),
+        }
+    }
+}
+
 /// Execute a command and return the output on stdout and stderr if successful
 #[cfg(not(test))]
 pub fn exec_cmd(cmd: &str, args: &[&str], time_limit: Duration) -> Option<CommandOutput> {
@@ -311,6 +323,16 @@ pub fn wrap_seq_for_shell(
 fn internal_exec_cmd(cmd: &str, args: &[&str], time_limit: Duration) -> Option<CommandOutput> {
     log::trace!("Executing command {:?} with args {:?}", cmd, args);
 
+    let cache_dir = std::env::var_os("STARSHIP_CACHE")
+        .map(PathBuf::from)
+        .unwrap_or_else(|| {
+            dirs_next::home_dir()
+                .expect("Unable to find home directory")
+                .join(".cache/starship")
+        });
+    let mut cache = Cache::new(&cache_dir.join("bin-cache")).ok()?;
+    log::debug!("Cache initialized: {:?}", &cache_dir);
+
     let full_path = match which::which(cmd) {
         Ok(full_path) => {
             log::trace!("Using {:?} as {:?}", full_path, cmd);
@@ -324,7 +346,31 @@ fn internal_exec_cmd(cmd: &str, args: &[&str], time_limit: Duration) -> Option<C
 
     let start = Instant::now();
 
-    let process = match Command::new(full_path)
+    let full_command = format!("{} {}", cmd, args.join(" "));
+    if let Some(output) = cache.get(&full_path, &full_command) {
+        log::info!("Retreived {:?} from cache: {:?}", full_command, output);
+
+        let stdout_string = match String::from_utf8(output.stdout.clone()) {
+            Ok(stdout) => stdout,
+            Err(error) => {
+                log::warn!("Unable to decode stdout: {:?}", error);
+                return None;
+            }
+        };
+        let stderr_string = match String::from_utf8(output.stderr.clone()) {
+            Ok(stderr) => stderr,
+            Err(error) => {
+                log::warn!("Unable to decode stderr: {:?}", error);
+                return None;
+            }
+        };
+        return Some(CommandOutput {
+            stdout: stdout_string,
+            stderr: stderr_string,
+        });
+    }
+
+    let process = match Command::new(&full_path)
         .args(args)
         .stderr(Stdio::piped())
         .stdout(Stdio::piped())
@@ -340,14 +386,14 @@ fn internal_exec_cmd(cmd: &str, args: &[&str], time_limit: Duration) -> Option<C
 
     match process.with_output_timeout(time_limit).terminating().wait() {
         Ok(Some(output)) => {
-            let stdout_string = match String::from_utf8(output.stdout) {
+            let stdout_string = match String::from_utf8(output.stdout.clone()) {
                 Ok(stdout) => stdout,
                 Err(error) => {
                     log::warn!("Unable to decode stdout: {:?}", error);
                     return None;
                 }
             };
-            let stderr_string = match String::from_utf8(output.stderr) {
+            let stderr_string = match String::from_utf8(output.stderr.clone()) {
                 Ok(stderr) => stderr,
                 Err(error) => {
                     log::warn!("Unable to decode stderr: {:?}", error);
@@ -355,6 +401,22 @@ fn internal_exec_cmd(cmd: &str, args: &[&str], time_limit: Duration) -> Option<C
                 }
             };
 
+            cache.set(
+                &full_path,
+                &full_command,
+                CachedOutput {
+                    stdout: output.stdout,
+                    stderr: output.stderr,
+                    status: output
+                        .status
+                        .code()
+                        .map(|i| i.try_into().unwrap_or_default()),
+                },
+            );
+            cache.write().unwrap_or_else(|e| {
+                log::warn!("Unable to write to binary cache: {}", e);
+            });
+
             log::trace!(
                 "stdout: {:?}, stderr: {:?}, exit code: \"{:?}\", took {:?}",
                 stdout_string,