1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
|
use std::ffi::OsString;
use std::io::{self, Write};
use std::path::PathBuf;
use std::process::{Command, exit, Stdio};
use std::time::{Duration, Instant};
use anyhow::{Context, Result};
use clap::Parser;
use bkt::{CommandDesc, Bkt};
// Re-invokes bkt with --force and then discards the subprocess, causing the cache
// to be refreshed asynchronously.
fn force_update_async() -> Result<()> {
let mut args = std::env::args_os();
let arg0 = args.next().expect("Must always be a 0th argument");
let mut command = match std::env::current_exe() {
Ok(path) => Command::new(path),
Err(_) => Command::new(arg0),
};
// Discard stdout/err so the calling process doesn't wait for them to close.
// Intentionally drop the returned Child; after this process exits the
// child process will continue running in the background.
command.arg("--force").args(args.filter(|a| a != "--warm"))
.stdout(Stdio::null()).stderr(Stdio::null())
.spawn().context("Failed to start background process")?;
Ok(())
}
// Runs bkt after main() handles flag parsing
fn run(cli: Cli) -> Result<i32> {
let ttl: Duration = cli.ttl.into();
let stale: Option<Duration> = cli.stale.map(Into::into);
assert!(!ttl.is_zero(), "--ttl cannot be zero");
if let Some(stale) = stale {
assert!(!stale.is_zero(), "--stale cannot be zero");
assert!(stale < ttl, "--stale must be less than --ttl");
}
let mut bkt = match cli.cache_dir {
Some(cache_dir) => Bkt::create(cache_dir)?,
None => Bkt::in_tmp()?,
};
if let Some(scope) = cli.scope {
bkt = bkt.scoped(scope);
}
let mut command = CommandDesc::new(cli.command);
if cli.cwd {
command = command.with_cwd();
}
let envs = cli.env;
if !envs.is_empty() {
command = command.with_envs(&envs);
}
let files = cli.modtime;
if !files.is_empty() {
command = command.with_modtimes(&files);
}
if cli.discard_failures {
command = command.with_discard_failures(true);
}
if cli.warm && !cli.force {
force_update_async()?;
return Ok(0);
}
let invocation = if cli.force {
bkt.refresh(&command, ttl)?.0
} else {
let (invocation, status) = bkt.retrieve(&command, ttl)?;
if let Some(stale) = stale {
if let bkt::CacheStatus::Hit(cached_at) = status {
if (Instant::now() - cached_at) > stale {
force_update_async()?;
}
}
}
invocation
};
// BrokenPipe errors are uninteresting for command line applications; just stop writing to that
// descriptor and, if appropriate, exit. Rust doesn't have good support for this presently, see
// https://github.com/rust-lang/rust/issues/46016
fn disregard_broken_pipe(result: std::io::Result<()>) -> std::io::Result<()> {
use std::io::ErrorKind::*;
if let Err(e) = &result {
if let BrokenPipe = e.kind() {
return Ok(());
}
}
result
}
disregard_broken_pipe(io::stdout().write_all(invocation.stdout()))
.context("error writing to stdout")?;
disregard_broken_pipe(io::stderr().write_all(invocation.stderr()))
.context("error writing to stderr")?;
Ok(invocation.exit_code())
}
#[derive(Debug, Parser)]
#[command(about, version)]
struct Cli {
/// The command to run
#[arg(required = true, last = true)]
command: Vec<OsString>,
/// Duration the cached result will be valid for
#[arg(long, value_name = "DURATION", default_value = "60s", visible_alias = "time-to-live", env = "BKT_TTL")]
ttl: humantime::Duration,
/// Duration after which the result will be asynchronously refreshed
#[arg(long, value_name = "DURATION", conflicts_with = "warm")]
stale: Option<humantime::Duration>,
/// Asynchronously execute and cache the given command, even if it's already cached
#[arg(long)]
warm: bool,
/// Execute and cache the given command, even if it's already cached
#[arg(long, conflicts_with = "warm")]
force: bool,
/// Includes the current working directory in the cache key,
/// so that the same command run in different directories caches separately
#[arg(long, visible_alias = "use-working-dir")]
cwd: bool,
/// Includes the given environment variable in the cache key,
/// so that the same command run with different values for the given variables caches separately
#[arg(long, value_name = "NAME", visible_alias = "use-environment")]
env: Vec<OsString>,
/// Includes the last modification time of the given file(s) in the cache key,
/// so that the same command run with different modtimes for the given files caches separately
#[arg(long, value_name = "FILE", visible_alias = "use-file-modtime")]
modtime: Vec<OsString>,
/// Don't cache invocations that fail (non-zero exit code).
/// USE CAUTION when passing this flag, as unexpected failures can lead to a spike in invocations
/// which can exacerbate ongoing issues, effectively a DDoS.
#[arg(long)]
discard_failures: bool,
/// If set, all cached data will be scoped to this value,
/// preventing collisions with commands cached with different scopes
#[arg(long, value_name = "NAME", env = "BKT_SCOPE")]
scope: Option<String>,
/// The directory under which to persist cached invocations;
/// defaults to the system's temp directory.
/// Setting this to a directory backed by RAM or an SSD, such as a tmpfs partition,
/// will significantly reduce caching overhead.
#[arg(long, value_name = "DIR", env = "BKT_CACHE_DIR")]
cache_dir: Option<PathBuf>,
}
fn main() {
let cli = Cli::parse();
match run(cli) {
Ok(code) => exit(code),
Err(msg) => {
eprintln!("bkt: {:#}", msg);
exit(127);
}
}
}
|
