Flesh out exec_cmd with cache

This commit is contained in:
Matan Kushner 2021-07-04 10:17:47 -04:00
parent 23539446b0
commit ad4705b7e9
No known key found for this signature in database
GPG Key ID: BFF10DF8EAA776DD
9 changed files with 530 additions and 51 deletions

56
Cargo.lock generated
View File

@ -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"
@ -394,6 +396,15 @@ dependencies = [
"winapi",
]
[[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"
@ -404,6 +415,17 @@ dependencies = [
"dirs-sys-next",
]
[[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"
@ -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",
@ -1733,6 +1756,17 @@ dependencies = [
"yaml-rust",
]
[[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"
@ -1856,6 +1890,26 @@ dependencies = [
"unicode-width",
]
[[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"

View File

@ -28,8 +28,10 @@ http = ["attohttpc"]
# Vendor OpenSSL, use this if you have trouble cross-compiling starship
tls-vendored = ["native-tls/vendored"]
[dependencies]
starship_cache = { version = "0.1.0", path = "crates/starship_cache" }
starship_module_config_derive = { version = "0.2.1", path = "crates/starship_module_config_derive" }
clap = "2.33.3"
ansi_term = "0.12.1"
dirs-next = "2.0.0"
@ -49,7 +51,6 @@ 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" }
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/*"]

View File

@ -0,0 +1,16 @@
[package]
name = "starship_cache"
version = "0.1.0"
authors = ["Starship Contributors"]
description = "Intelligent caching for Starship"
edition = "2018"
license = "ISC"
[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"

View File

@ -0,0 +1,20 @@
use std::io;
#[non_exhaustive]
#[derive(thiserror::Error, Debug)]
pub enum Error {
#[error("cannot open cache file")]
OpenFile(#[source] io::Error),
#[error("cannot create cache dir")]
CreateCacheDir(#[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),
}

View File

@ -0,0 +1,317 @@
//! 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},
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 create_or_parse<P: AsRef<Path>>(cache_dir: P) -> Result<Self, Error> {
let cache_dir = cache_dir.as_ref();
fs::create_dir_all(&cache_dir)
.map_err(Error::CreateCacheDir)?;
let cache_file = cache_dir.join("bin-cache");
let mut file = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.open(&cache_file)
.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: cache_file,
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, Default, Debug, PartialEq)]
pub struct CachedOutput {
pub stdout: String,
pub stderr: String,
pub status: Option<i32>,
}
impl CachedOutput {
pub fn success(&self) -> bool {
self.status == Some(0)
}
}
#[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>>;
// Convenience
impl From<&str> for CachedOutput {
fn from(stdout: &str) -> Self {
Self {
stdout: stdout.to_owned(),
..Default::default()
}
}
}
#[test]
fn empty_cache_file_is_created() -> Result {
let dir = tempdir()?;
let cache_path = Path::join(dir.path(), "bin-cache");
let cache = Cache::create_or_parse(&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::create_or_parse(&cache_path)?;
// Create "node" binary
let bin_path = dir.path().join("node");
File::create(&bin_path)?;
// Populate cache with "node" output
let version = "v14.16.0";
cache.set(&bin_path, "node --version", version);
cache.write()?;
// Retreive cached output
let mut new_cache = Cache::create_or_parse(&cache_path)?;
let actual = new_cache.get(&bin_path, "node --version").unwrap();
let expected = CachedOutput::from(version);
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::create_or_parse(&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::create_or_parse(&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::create_or_parse(&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::create_or_parse(&cache_path)?;
let actual = new_cache.get(&bin_path, "node --version");
assert_eq!(None, actual);
Ok(())
}
}

View File

@ -7,6 +7,7 @@ use clap::ArgMatches;
use dirs_next::home_dir;
use git2::{ErrorCode::UnbornBranch, Repository, RepositoryState};
use once_cell::sync::OnceCell;
use starship_cache::Cache;
use std::collections::{HashMap, HashSet};
use std::env;
use std::ffi::OsString;
@ -44,17 +45,21 @@ pub struct Context<'a> {
/// A HashMap of environment variable mocks
#[cfg(test)]
pub env: HashMap<&'a str, String>,
pub env_mocks: HashMap<&'a str, String>,
/// A HashMap of command mocks
#[cfg(test)]
pub cmd: HashMap<&'a str, Option<CommandOutput>>,
pub cmd_mocks: HashMap<&'a str, Option<CommandOutput>>,
/// A placeholder for a mock battery provider
#[cfg(feature = "battery")]
pub battery_info_provider: &'a (dyn crate::modules::BatteryInfoProvider + Send + Sync),
/// Timeout for the execution of commands
cmd_timeout: Duration,
/// An instance of Starship's binary output cache
cmd_cache: Cache
}
impl<'a> Context<'a> {
@ -113,6 +118,15 @@ impl<'a> Context<'a> {
let cmd_timeout = Duration::from_millis(config.get_root_config().command_timeout);
let cache_dir = 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 cmd_cache = Cache::create_or_parse(cache_dir).unwrap();
Context {
config,
properties,
@ -122,12 +136,13 @@ impl<'a> Context<'a> {
repo: OnceCell::new(),
shell,
#[cfg(test)]
env: HashMap::new(),
env_mocks: HashMap::new(),
#[cfg(test)]
cmd: HashMap::new(),
cmd_mocks: HashMap::new(),
#[cfg(feature = "battery")]
battery_info_provider: &crate::modules::BatteryInfoProviderImpl,
cmd_timeout,
cmd_cache
}
}
@ -143,7 +158,7 @@ impl<'a> Context<'a> {
// Retrives a environment variable from the os or from a table if in testing mode
#[cfg(test)]
pub fn get_env<K: AsRef<str>>(&self, key: K) -> Option<String> {
self.env.get(key.as_ref()).map(|val| val.to_string())
self.env_mocks.get(key.as_ref()).map(|val| val.to_string())
}
#[cfg(not(test))]
@ -155,7 +170,7 @@ impl<'a> Context<'a> {
// Retrives a environment variable from the os or from a table if in testing mode (os version)
#[cfg(test)]
pub fn get_env_os<K: AsRef<str>>(&self, key: K) -> Option<OsString> {
self.env.get(key.as_ref()).map(OsString::from)
self.env_mocks.get(key.as_ref()).map(OsString::from)
}
#[cfg(not(test))]
@ -264,20 +279,58 @@ impl<'a> Context<'a> {
self.properties.get("cmd_duration")?.parse::<u128>().ok()
}
/// Execute a command and return the output on stdout and stderr if successful
/// Execute a command and return the output on stdout and stderr if successful,
/// while caching the output and respecting the user's timeout configuration.
///
/// This method automatically caches successful commands to short-circuit
/// future calls. For a non-caching alternative, use [`Self::uncached_exec_cmd()`].
#[inline]
pub fn exec_cmd(&self, cmd: &str, args: &[&str]) -> Option<CommandOutput> {
pub fn exec_cmd(&mut self, cmd: &str, args: &[&str]) -> Option<CommandOutput> {
#[cfg(test)]
{
let command = match args.len() {
let full_command = match args.len() {
0 => cmd.to_owned(),
_ => format!("{} {}", cmd, args.join(" ")),
};
if let Some(output) = self.cmd.get(command.as_str()) {
if let Some(output) = self.cmd_mocks.get(full_command.as_str()) {
return output.clone();
}
}
log::trace!("Executing command {:?} with args {:?}", cmd, args);
let full_path = match which::which(cmd) {
Ok(full_path) => {
log::trace!("Using {:?} as {:?}", full_path, cmd);
full_path
}
Err(error) => {
log::trace!("Unable to find {:?} in PATH, {:?}", cmd, error);
return None;
}
};
let full_command = format!("{} {}", cmd, args.join(" "));
if let Some(output) = self.cmd_cache.get(&full_path, &full_command) {
log::info!("Retreived {:?} from cache: {:?}", full_command, output);
let output = CommandOutput::from(output);
return Some(output);
};
exec_cmd(cmd, args, self.cmd_timeout)
.map(|output| {
self.cmd_cache.set(&full_path, &full_command, &output);
output
})
}
/// Execute a command and return the output on stdout and stderr if successful,
/// while respecting the user's timeout configuration.
///
/// This method specifically doesn't cache its results. For an alternative that
/// caches for use by successive calls, use [`Self::exec_cmd()`].
pub fn uncached_exec_cmd(&self, cmd: &str, args: &[&str]) -> Option<CommandOutput> {
todo!()
}
}

View File

@ -317,10 +317,7 @@ mod tests {
File::create(dir.path().join("any.ml"))?.sync_all()?;
let actual = ModuleRenderer::new("ocaml")
.cmd(
"opam switch show --safe",
Some(CommandOutput::default()),
)
.cmd("opam switch show --safe", Some(CommandOutput::default()))
.path(dir.path())
.collect();
let expected = Some(format!("via {}", Color::Yellow.bold().paint("🐫 v4.10.0 ")));
@ -392,7 +389,7 @@ mod tests {
"opam switch show --safe",
Some(CommandOutput {
stdout: String::from("/path/to/my-project\n"),
..Default::default()
..Default::default()
}),
)
.path(dir.path())

View File

@ -79,13 +79,13 @@ impl<'a> ModuleRenderer<'a> {
/// Adds the variable to the env_mocks of the underlying context
pub fn env<V: Into<String>>(mut self, key: &'a str, val: V) -> Self {
self.context.env.insert(key, val.into());
self.context.env_mocks.insert(key, val.into());
self
}
/// Adds the command to the commandv_mocks of the underlying context
pub fn cmd(mut self, key: &'a str, val: Option<CommandOutput>) -> Self {
self.context.cmd.insert(key, val);
self.context.cmd_mocks.insert(key, val);
self
}

View File

@ -1,7 +1,9 @@
use process_control::{ChildExt, Timeout};
use starship_cache::CachedOutput;
use std::convert::TryFrom;
use std::fmt::Debug;
use std::fs::read_to_string;
use std::io::Result;
use std::io;
use std::path::Path;
use std::process::{Command, Stdio};
use std::time::{Duration, Instant};
@ -9,7 +11,7 @@ use std::time::{Duration, Instant};
use crate::context::Shell;
/// Return the string contents of a file
pub fn read_file<P: AsRef<Path> + Debug>(file_name: P) -> Result<String> {
pub fn read_file<P: AsRef<Path> + Debug>(file_name: P) -> io::Result<String> {
log::trace!("Trying to read from {:?}", file_name);
let result = read_to_string(file_name);
@ -30,6 +32,43 @@ pub struct CommandOutput {
pub status: i64,
}
impl TryFrom<process_control::Output> for CommandOutput {
type Error = String;
fn try_from(output: process_control::Output) -> Result<Self, Self::Error> {
let stdout = String::from_utf8(output.stdout)
.map_err(|err| format!("Unable to decode stdout: {:?}", err))?;
let stderr = String::from_utf8(output.stderr)
.map_err(|err| format!("Unable to decode stderr: {:?}", err))?;
Ok(Self {
stdout,
stderr,
status: output.status.code().unwrap_or_default(),
})
}
}
impl From<&CachedOutput> for CommandOutput {
fn from(output: &CachedOutput) -> Self {
Self {
stdout: output.stdout.to_owned(),
stderr: output.stderr.to_owned(),
status: output.status.unwrap_or_default() as i64
}
}
}
impl Into<CachedOutput> for &CommandOutput {
fn into(self) -> CachedOutput {
CachedOutput {
stdout: self.stdout.clone(),
stderr: self.stdout.clone(),
status: Some(i32::try_from(self.status).unwrap_or_default())
}
}
}
/// 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> {
@ -305,7 +344,7 @@ 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 start = Instant::now();
let full_path = match which::which(cmd) {
Ok(full_path) => {
@ -318,8 +357,6 @@ fn internal_exec_cmd(cmd: &str, args: &[&str], time_limit: Duration) -> Option<C
}
};
let start = Instant::now();
let process = match Command::new(full_path)
.args(args)
.stderr(Stdio::piped())
@ -336,38 +373,19 @@ 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) {
Ok(stdout) => stdout,
Err(error) => {
log::warn!("Unable to decode stdout: {:?}", error);
return None;
}
};
let stderr_string = match String::from_utf8(output.stderr) {
Ok(stderr) => stderr,
Err(error) => {
log::warn!("Unable to decode stderr: {:?}", error);
return None;
}
};
let output = CommandOutput::try_from(output)
.map_err(|err| log::warn!("{}", err))
.ok()?;
log::trace!(
"stdout: {:?}, stderr: {:?}, exit code: \"{:?}\", took {:?}",
stdout_string,
stderr_string,
output.status.code(),
output.stdout,
output.stderr,
output.status,
start.elapsed()
);
if !output.status.success() {
return None;
}
Some(CommandOutput {
stdout: stdout_string,
stderr: stderr_string,
status: output.status.code().unwrap_or_default()
})
Some(output)
}
Ok(None) => {
log::warn!("Executing command {:?} timed out.", cmd);