use crate::materialize::{
CachedCompilation, CachedOutput, Materialization, StagedOutputs, apply_file_mode,
denormalize_output_text, executable_mode_matches, exit_code, file_mode, find_blobs,
normalize_output_text, persist_outputs, read_canonical_blob, read_verified_blob,
record_action_hit, record_verification, replay_bytes, resolve_executable,
stage_verified_cached_output, staging_directory, validate_file_mode,
};
use crate::{session, util::workspace_root};
use eyre::{Context, Result, bail};
use mbx_cache_core::{
ActionPrediction, AgentRequest, AgentResponse, CacheDigest, CacheDirectory, CacheFileNode,
FileDigestScope, FileIdentity, RecordedFileDigest, RemoteActionResult, RestoreStats,
RustcMetadata, canonical_json,
};
use mbx_cache_rustc::{
ActionContext, BypassReason, CompilerIdentity, DiscoveredInputs, LinkerIdentity, ParseOptions,
PathMapping, RustcAction, RustcDepInfo, RustcInputPrediction, RustcInvocation, RustcOutputs,
normalize_mapped_path,
};
use serde::{Deserialize, Serialize};
use std::borrow::Cow;
use std::collections::{BTreeMap, BTreeSet};
use std::ffi::{OsStr, OsString};
use std::path::{Path, PathBuf};
use std::process::{Command, ExitCode, Output};
use std::sync::OnceLock;
use std::time::{Instant, SystemTime};
#[derive(Clone, Debug, Default)]
struct CompileTiming {
crate_name: String,
duration_ns: u64,
}
const HOT_STREAK_THRESHOLD: u32 = 3;
const CHURN_STATE_VERSION: u8 = 1;
const INCREMENTAL_DIR_MAX_BYTES: u64 = 1024 * 1024 * 1024;
#[derive(Clone, Debug, Default)]
struct LearnedPlan {
hot: bool,
streak: u32,
directory: Option<PathBuf>,
record: Option<(PathBuf, CacheDigest)>,
}
impl LearnedPlan {
fn resolved(mut self, invocation: &CacheDigest) -> Self {
if self.hot {
self.directory = incremental_directory(invocation);
}
self
}
fn record_compiled(&self) {
let Some((path, sources)) = &self.record else {
return;
};
if let Err(error) = write_churn_state(path, sources, self.streak) {
eprintln!("mbx[warning]: churn was not recorded for this crate: {error:#}");
}
}
fn engaged(&self) -> bool {
self.directory.is_some()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct ChurnState {
version: u8,
sources: String,
streak: u32,
}
struct Compilation<'a> {
rustc: &'a OsStr,
invocation: &'a RustcInvocation,
working_dir: &'a Path,
portable: &'a Portable,
linker: Option<LinkerIdentity>,
}
pub(crate) fn compile(rustc: &OsStr, arguments: &[OsString]) -> Result<ExitCode> {
let working_dir = std::env::current_dir()?;
let options = ParseOptions::caching_native_links(session::cache_links_requested());
let arguments = with_oso_prefix(arguments, session::cache_links_requested());
let arguments = arguments.as_ref();
let initial_invocation = RustcInvocation::parse_with(arguments, options)?;
let initial_outputs = initial_invocation.outputs(&working_dir)?;
let portable = Portable::detect(
&working_dir,
Some(&initial_outputs.directory),
initial_invocation.target(),
);
let arguments = portable.applied_to(arguments);
let invocation = RustcInvocation::parse_with(&arguments, options)?;
let outputs = invocation.outputs(&working_dir)?;
let verify = session::verify_requested();
let learned_enabled = session::learned_incremental_requested() && !verify;
let mut verification = None;
let mut action_lookup_attempted = false;
let mut learned = LearnedPlan::default();
let compilation = Compilation {
rustc,
invocation: &invocation,
working_dir: &working_dir,
portable: &portable,
linker: linker_for(&invocation)?,
};
if outputs.dep_info.is_file()
&& let Ok((candidates, discovered)) =
action_from_current_dep_info(&compilation, &outputs.dep_info)
{
action_lookup_attempted = true;
match restore_candidates(
&candidates,
&outputs,
&discovered,
!verify,
&portable.mappings,
) {
Ok(Some((action, mut cached))) => {
match refresh_prediction(&compilation, &action, &discovered) {
Ok(timing) => {
cached.restore.avoided_compiler_duration_ns = timing.duration_ns;
}
Err(error) => {
eprintln!("mbx[warning]: compiler timing was not refreshed: {error:#}");
}
}
if verify {
verification = Some(cached);
} else {
record_action_hit(&action, cached.restore, invocation.crate_name());
let _ = replay_bytes(&cached.stdout, &cached.stderr);
return Ok(ExitCode::SUCCESS);
}
}
Ok(None) => {
learned = plan_learned_reuse(&compilation, &discovered, learned_enabled);
}
Err(error) => {
eprintln!("mbx[warning]: result was not restored: {error:#}");
}
}
}
let mut prediction_missing = false;
if !action_lookup_attempted {
match restore_predicted_result(
&compilation,
&outputs,
!verify,
&mut action_lookup_attempted,
learned_enabled,
&mut learned,
) {
Ok(Some(cached)) => {
if verify {
verification = Some(cached);
} else {
let _ = replay_bytes(&cached.stdout, &cached.stderr);
return Ok(ExitCode::SUCCESS);
}
}
Ok(None) => {
prediction_missing = !action_lookup_attempted;
}
Err(error) => {
eprintln!("mbx[warning]: prediction was not restored: {error:#}");
}
}
}
let flight = if verify {
None
} else {
join_flight(&compilation)
};
if let Some(flight) = &flight
&& (flight.flight.waited() || !action_lookup_attempted)
&& let Some(payload) = flight.flight.inherited()
{
match restore_flight_prediction(
&compilation,
&outputs,
&flight.invocation,
payload,
&mut action_lookup_attempted,
learned_enabled,
&mut learned,
) {
Ok(Some(cached)) => {
let _ = replay_bytes(&cached.stdout, &cached.stderr);
return Ok(ExitCode::SUCCESS);
}
Ok(None) => {}
Err(error) => {
eprintln!("mbx[warning]: a flight prediction was not restored: {error:#}");
}
}
}
if prediction_missing && !action_lookup_attempted {
session::record_unconsulted();
}
let demand =
crate::scheduler::Demand::new(invocation.crate_name(), invocation.links_natively());
let permit = crate::scheduler::pool().and_then(|pool| pool.admit(&demand));
let compilation_started = SystemTime::now();
let compiler_timer = Instant::now();
let mut command = Command::new(rustc);
command.args(&arguments).current_dir(&working_dir);
if let Some(directory) = learned.directory.as_deref() {
let mut flag = OsString::from("-Cincremental=");
flag.push(directory);
command.arg(flag);
}
let output = command.output().wrap_err("failed to execute rustc")?;
drop(permit);
crate::scheduler::record_compiler_memory(&demand, &output.status);
let timing = CompileTiming {
crate_name: invocation.crate_name().to_string(),
duration_ns: compiler_timer
.elapsed()
.as_nanos()
.try_into()
.unwrap_or(u64::MAX),
};
session::record_compiler_invocation(
if verification.is_some() {
"verification"
} else if learned.engaged() {
"incremental"
} else if action_lookup_attempted {
"miss"
} else {
"unconsulted"
},
Some(&timing.crate_name),
timing.duration_ns,
);
let _ = replay_output(&output);
if let Some(cached) = verification {
let divergence = verification_divergence(&cached, &output);
record_verification(divergence.is_none(), cached.restore);
if let Some(divergence) = divergence {
eprintln!(
"mbx[warning]: shadow verification diverged from cached output: {divergence}"
);
}
return Ok(exit_code(output.status));
}
if output.status.success() {
learned.record_compiled();
let publication: Result<()> = (|| {
let (candidates, discovered) = action_from_dep_info(&compilation, &outputs.dep_info)?;
discovered.verify_not_modified_since(compilation_started)?;
discovered.verify()?;
let action = if learned.engaged() {
&candidates.literal
} else {
publish_result(
&candidates,
&portable,
&outputs,
&output,
&portable.mappings,
)?
};
record_prediction(
&compilation,
&action.digest,
&discovered,
&timing,
flight
.as_ref()
.filter(|_| !learned.engaged())
.map(|flight| &flight.flight),
);
Ok(())
})();
if let Err(error) = publication {
eprintln!("mbx[warning]: result was not stored: {error:#}");
}
}
Ok(exit_code(output.status))
}
fn learned_plan(
recorded: Option<&ChurnState>,
sources: &CacheDigest,
enabled: bool,
) -> LearnedPlan {
let streak = match recorded {
Some(recorded) if recorded.sources == sources.key() => 0,
Some(recorded) => recorded.streak.saturating_add(1).min(HOT_STREAK_THRESHOLD),
None => 0,
};
LearnedPlan {
hot: enabled && streak >= HOT_STREAK_THRESHOLD,
streak,
directory: None,
record: None,
}
}
fn incremental_directory(invocation: &CacheDigest) -> Option<PathBuf> {
let target_dir = std::env::var_os(session::TARGET_DIR_ENV)?;
let key = invocation.key();
let shard = key.get(..16)?;
let directory = PathBuf::from(target_dir)
.join("mbx-incremental")
.join(shard);
match prepare_incremental_directory(&directory) {
Ok(()) => Some(directory),
Err(error) => {
eprintln!("mbx[warning]: incremental state was not prepared: {error:#}");
None
}
}
}
fn prepare_incremental_directory(directory: &Path) -> Result<()> {
if directory_bytes(directory) > INCREMENTAL_DIR_MAX_BYTES {
std::fs::remove_dir_all(directory)
.wrap_err("failed to discard oversized incremental state")?;
}
std::fs::create_dir_all(directory).wrap_err("failed to create the incremental directory")
}
fn directory_bytes(directory: &Path) -> u64 {
let mut total = 0;
let mut pending = vec![directory.to_path_buf()];
while let Some(next) = pending.pop() {
let Ok(listing) = std::fs::read_dir(&next) else {
continue;
};
for entry in listing.flatten() {
let Ok(metadata) = entry.metadata() else {
continue;
};
if metadata.is_dir() {
pending.push(entry.path());
} else if metadata.is_file() {
total += metadata.len();
}
}
}
total
}
fn plan_learned_reuse(
compilation: &Compilation<'_>,
discovered: &DiscoveredInputs,
enabled: bool,
) -> LearnedPlan {
let planned = (|| {
let sources = compilation.invocation.source_fingerprint(discovered);
let context = base_action_context(
compilation.rustc,
compilation.working_dir,
compilation.portable,
)?;
let unit = compilation.invocation.invocation_digest(&context)?;
let Some(state_path) = churn_state_path(&unit) else {
return Result::<LearnedPlan>::Ok(LearnedPlan::default());
};
let plan = learned_plan(read_churn_state(&state_path).as_ref(), &sources, enabled);
Ok(LearnedPlan {
record: Some((state_path, sources)),
..plan
}
.resolved(&unit))
})();
match planned {
Ok(plan) => plan,
Err(error) => {
eprintln!("mbx[warning]: churn was not tracked for this crate: {error:#}");
LearnedPlan::default()
}
}
}
fn churn_state_path(unit: &CacheDigest) -> Option<PathBuf> {
let target_dir = std::env::var_os(session::TARGET_DIR_ENV)?;
let key = unit.key();
let shard = key.get(..16)?;
Some(
PathBuf::from(target_dir)
.join("mbx-incremental")
.join(format!("{shard}.json")),
)
}
fn read_churn_state(path: &Path) -> Option<ChurnState> {
let bytes = std::fs::read(path).ok()?;
serde_json::from_slice::<ChurnState>(&bytes)
.ok()
.filter(|state| state.version == CHURN_STATE_VERSION)
}
fn write_churn_state(path: &Path, sources: &CacheDigest, streak: u32) -> Result<()> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.wrap_err("failed to create the incremental state directory")?;
}
let state = ChurnState {
version: CHURN_STATE_VERSION,
sources: sources.key(),
streak,
};
crate::util::write_atomic(path, &serde_json::to_vec(&state)?)
.wrap_err("failed to record what this crate last compiled")
}
fn restore_predicted_result(
compilation: &Compilation<'_>,
outputs: &RustcOutputs,
restore_outputs: bool,
action_lookup_attempted: &mut bool,
learned_enabled: bool,
learned: &mut LearnedPlan,
) -> Result<Option<CachedCompilation>> {
let Compilation {
invocation,
working_dir,
portable,
..
} = compilation;
let context = base_action_context(compilation.rustc, working_dir, portable)?;
let invocation_digest = invocation.invocation_digest(&context)?;
let task = prediction_task(&invocation_digest);
let responses = session::request_agent(&[AgentRequest::FindActionPrediction {
task,
invocation: invocation_digest.clone(),
}])?;
let Some(response) = responses.into_iter().next() else {
bail!("cache agent did not return an action prediction response");
};
let prediction = match response {
AgentResponse::ActionPrediction {
prediction: Some(prediction),
} => prediction,
AgentResponse::ActionPrediction { prediction: None } => {
return Ok(None);
}
AgentResponse::Error { message } => bail!(message),
_ => bail!("cache agent returned an unexpected action prediction response"),
};
if prediction.adapter != "rustc" || prediction.invocation != invocation_digest {
bail!("cache agent returned an incompatible rustc action prediction");
}
restore_prediction_payload(
compilation,
outputs,
context,
&invocation_digest,
&prediction.payload,
Some(&prediction.action),
restore_outputs,
action_lookup_attempted,
learned_enabled,
learned,
)
}
#[allow(clippy::too_many_arguments)]
fn restore_prediction_payload(
compilation: &Compilation<'_>,
outputs: &RustcOutputs,
mut context: ActionContext,
invocation_digest: &CacheDigest,
payload: &str,
recorded_action: Option<&CacheDigest>,
restore_outputs: bool,
action_lookup_attempted: &mut bool,
learned_enabled: bool,
learned: &mut LearnedPlan,
) -> Result<Option<CachedCompilation>> {
let Compilation {
invocation,
working_dir,
portable,
..
} = compilation;
let input_prediction: RustcInputPrediction = serde_json::from_str(payload)?;
if String::from_utf8(canonical_json(&input_prediction)?)? != payload {
bail!("the rustc action prediction is not canonical");
}
let discovered = input_prediction.discover(
working_dir,
&context.path_mappings,
session::file_digest_cache(),
)?;
discovered.clone().apply_to(&mut context)?;
let candidates = ActionCandidates::build(invocation, context, compilation.linker.clone())?;
*action_lookup_attempted = true;
let restored = restore_candidates(
&candidates,
outputs,
&discovered,
restore_outputs,
&portable.mappings,
)?;
match restored {
Some((action, mut cached)) => {
cached.restore.avoided_compiler_duration_ns = input_prediction.compiler_duration_ns;
if restore_outputs {
record_action_hit(&action, cached.restore, invocation.crate_name());
}
if recorded_action != Some(&action) {
record_prediction_value(
invocation_digest.clone(),
action,
payload.to_string(),
invocation.crate_name(),
);
}
Ok(Some(cached))
}
None => {
if learned.record.is_none() {
*learned = plan_learned_reuse(compilation, &discovered, learned_enabled);
}
Ok(None)
}
}
}
struct ActionCandidates {
portable: Option<RustcAction>,
literal: RustcAction,
}
impl ActionCandidates {
fn build(
invocation: &RustcInvocation,
context: ActionContext,
linker: Option<LinkerIdentity>,
) -> Result<Self> {
let applies = context
.portable_environment
.iter()
.any(|name| context.environment.contains_key(name));
let literal_context = ActionContext {
portable_environment: BTreeSet::new(),
..context.clone()
};
Ok(Self {
portable: applies
.then(|| invocation.action_linked_by(context, linker.clone()))
.transpose()?,
literal: invocation.action_linked_by(literal_context, linker)?,
})
}
fn publishable(&self, outputs_are_clean: bool) -> &RustcAction {
match &self.portable {
Some(action) if outputs_are_clean => action,
_ => &self.literal,
}
}
fn ordered(&self) -> impl Iterator<Item = &RustcAction> {
self.portable.iter().chain(std::iter::once(&self.literal))
}
}
fn restore_candidates(
candidates: &ActionCandidates,
outputs: &RustcOutputs,
discovered: &DiscoveredInputs,
restore_outputs: bool,
mappings: &[PathMapping],
) -> Result<Option<(CacheDigest, CachedCompilation)>> {
for action in candidates.ordered() {
if let Some(cached) =
restore_result(action, outputs, discovered, restore_outputs, mappings)?
{
return Ok(Some((action.digest.clone(), cached)));
}
}
Ok(None)
}
fn action_from_dep_info(
compilation: &Compilation<'_>,
dep_info: &Path,
) -> Result<(ActionCandidates, DiscoveredInputs)> {
let dep_info = RustcDepInfo::read(dep_info)?;
action_from_parsed_dep_info(compilation, &dep_info)
}
fn action_from_current_dep_info(
compilation: &Compilation<'_>,
dep_info: &Path,
) -> Result<(ActionCandidates, DiscoveredInputs)> {
let dep_info = RustcDepInfo::read(dep_info)?;
verify_environment(&dep_info.environment)?;
action_from_parsed_dep_info(compilation, &dep_info)
}
fn action_from_parsed_dep_info(
compilation: &Compilation<'_>,
dep_info: &RustcDepInfo,
) -> Result<(ActionCandidates, DiscoveredInputs)> {
let Compilation {
invocation,
working_dir,
portable,
..
} = compilation;
let discovered = invocation.discover_inputs_with_mappings(
dep_info,
working_dir,
&portable.mappings,
session::file_digest_cache(),
)?;
let mut context = base_action_context(compilation.rustc, working_dir, portable)?;
discovered.clone().apply_to(&mut context)?;
let candidates = ActionCandidates::build(invocation, context, compilation.linker.clone())?;
Ok((candidates, discovered))
}
fn base_action_context(
rustc: &OsStr,
working_dir: &Path,
portable: &Portable,
) -> Result<ActionContext> {
Ok(ActionContext {
compiler: compiler_identity(rustc)?,
working_dir: working_dir.to_path_buf(),
path_mappings: portable.mappings.clone(),
environment: BTreeMap::new(),
portable_environment: portable.names.clone(),
inputs: Vec::new(),
})
}
fn with_oso_prefix(arguments: &[OsString], cache_links: bool) -> Cow<'_, [OsString]> {
if !cfg!(target_os = "macos") || !cache_links {
return Cow::Borrowed(arguments);
}
let mut out_dir = None;
for (index, argument) in arguments.iter().enumerate() {
let Some(argument) = argument.to_str() else {
continue;
};
if argument.contains("-oso_prefix,") {
return Cow::Borrowed(arguments);
}
if argument == "--target" || argument.starts_with("--target=") {
return Cow::Borrowed(arguments);
}
if argument == "--out-dir" {
out_dir = arguments.get(index + 1).map(PathBuf::from);
} else if let Some(value) = argument.strip_prefix("--out-dir=") {
out_dir = Some(PathBuf::from(value));
}
}
let Some(out_dir) = out_dir.filter(|out_dir| out_dir.is_absolute()) else {
return Cow::Borrowed(arguments);
};
let mut extended = arguments.to_vec();
extended.push(format!("-Clink-arg=-Wl,-oso_prefix,{}/", out_dir.display()).into());
Cow::Owned(extended)
}
fn linker_for(invocation: &RustcInvocation) -> Result<Option<LinkerIdentity>> {
if !invocation.links_natively() {
return Ok(None);
}
match crate::linker::identity() {
Ok(identity) => Ok(Some(identity)),
Err(error) => Err(BypassReason::UnportableNativeLink(format!(
"the linker could not be identified: {error:#}"
))
.into()),
}
}
fn record_prediction(
compilation: &Compilation<'_>,
action: &CacheDigest,
discovered: &DiscoveredInputs,
timing: &CompileTiming,
flight: Option<&crate::scheduler::Flight>,
) {
let result = (|| {
let invocation = compilation.invocation;
let context = base_action_context(
compilation.rustc,
compilation.working_dir,
compilation.portable,
)?;
let invocation_digest = invocation.invocation_digest(&context)?;
let mut prediction = invocation.prediction(&context, discovered)?;
prediction.version = prediction.version.max(2);
prediction.compiler_duration_ns = timing.duration_ns;
prediction.crate_name.clone_from(&timing.crate_name);
let payload = String::from_utf8(canonical_json(&prediction)?)?;
if let Some(flight) = flight {
flight.leave(&payload);
}
record_prediction_value(
invocation_digest,
action.clone(),
payload,
invocation.crate_name(),
);
Result::<()>::Ok(())
})();
if let Err(error) = result {
warn_prediction_not_recorded(compilation.invocation.crate_name(), &error);
}
}
struct InvocationFlight {
flight: crate::scheduler::Flight,
invocation: CacheDigest,
}
fn join_flight(compilation: &Compilation<'_>) -> Option<InvocationFlight> {
let invocation = (|| -> Result<CacheDigest> {
let context = base_action_context(
compilation.rustc,
compilation.working_dir,
compilation.portable,
)?;
Ok(compilation.invocation.invocation_digest(&context)?)
})()
.ok()?;
let flight = crate::scheduler::flight("rustc", &invocation.hash)?;
Some(InvocationFlight { flight, invocation })
}
fn restore_flight_prediction(
compilation: &Compilation<'_>,
outputs: &RustcOutputs,
invocation_digest: &CacheDigest,
payload: &str,
action_lookup_attempted: &mut bool,
learned_enabled: bool,
learned: &mut LearnedPlan,
) -> Result<Option<CachedCompilation>> {
let context = base_action_context(
compilation.rustc,
compilation.working_dir,
compilation.portable,
)?;
restore_prediction_payload(
compilation,
outputs,
context,
invocation_digest,
payload,
None,
true,
action_lookup_attempted,
learned_enabled,
learned,
)
}
fn refresh_prediction(
compilation: &Compilation<'_>,
action: &CacheDigest,
discovered: &DiscoveredInputs,
) -> Result<CompileTiming> {
let context = base_action_context(
compilation.rustc,
compilation.working_dir,
compilation.portable,
)?;
let invocation_digest = compilation.invocation.invocation_digest(&context)?;
let task = prediction_task(&invocation_digest);
let responses = session::request_agent(&[AgentRequest::FindActionPrediction {
task,
invocation: invocation_digest.clone(),
}])?;
let Some(response) = responses.into_iter().next() else {
bail!("cache agent did not return an action prediction response");
};
let stored = match response {
AgentResponse::ActionPrediction { prediction } => prediction,
AgentResponse::Error { message } => bail!(message),
_ => bail!("cache agent returned an unexpected action prediction response"),
};
let timing = match &stored {
Some(stored) => decode_prediction_timing(stored, &invocation_digest)?,
None => CompileTiming::default(),
};
let recorded = (|| {
let mut prediction = compilation.invocation.prediction(&context, discovered)?;
prediction.version = prediction.version.max(2);
prediction.compiler_duration_ns = timing.duration_ns;
prediction.crate_name.clone_from(&timing.crate_name);
let payload = String::from_utf8(canonical_json(&prediction)?)?;
let unchanged = stored
.as_ref()
.is_some_and(|stored| stored.action == *action && stored.payload == payload);
if !unchanged {
record_prediction_value(
invocation_digest,
action.clone(),
payload,
compilation.invocation.crate_name(),
);
}
Result::<()>::Ok(())
})();
if let Err(error) = recorded {
warn_prediction_not_recorded(compilation.invocation.crate_name(), &error);
}
Ok(timing)
}
fn decode_prediction_timing(
prediction: &ActionPrediction,
invocation: &CacheDigest,
) -> Result<CompileTiming> {
if prediction.adapter != "rustc" || prediction.invocation != *invocation {
bail!("cache agent returned an incompatible rustc timing prediction");
}
let timing: RustcInputPrediction = serde_json::from_str(&prediction.payload)?;
if !matches!(timing.version, 1..=3)
|| timing.crate_name.len() > 256
|| timing.crate_name.contains(['\0', '\n', '\r'])
|| String::from_utf8(canonical_json(&timing)?)? != prediction.payload
{
bail!("cache agent returned an invalid rustc timing prediction");
}
Ok(CompileTiming {
crate_name: timing.crate_name,
duration_ns: timing.compiler_duration_ns,
})
}
fn record_prediction_value(
invocation: CacheDigest,
action: CacheDigest,
payload: String,
crate_name: &str,
) {
let result = (|| {
let task = prediction_task(&invocation);
let responses = session::request_agent(&[AgentRequest::RecordActionPrediction {
task,
prediction: ActionPrediction {
invocation,
action,
adapter: "rustc".into(),
payload,
},
}])?;
match responses.into_iter().next() {
Some(AgentResponse::ActionPredictionRecorded) => Ok(()),
Some(AgentResponse::Error { message }) => bail!(message),
_ => bail!("cache agent returned an unexpected prediction response"),
}
})();
if let Err(error) = result {
warn_prediction_not_recorded(crate_name, &error);
}
}
fn warn_prediction_not_recorded(crate_name: &str, error: &eyre::Report) {
eprintln!("mbx[warning]: action prediction for {crate_name} was not recorded: {error:#}");
}
fn prediction_task(invocation: &CacheDigest) -> String {
std::env::var(session::BUILD_ENV).unwrap_or_else(|_| {
let shard = invocation.hash.get(..2).unwrap_or(&invocation.hash);
CacheDigest::blake3(format!("standalone-predictions-v1\0{shard}").as_bytes()).hash
})
}
fn verify_environment(environment: &BTreeMap<String, Option<String>>) -> Result<()> {
for (name, expected) in environment {
let actual = std::env::var_os(name)
.map(|value| {
value.into_string().map_err(|_| {
eyre::eyre!("compiler environment input is not valid UTF-8: {name}")
})
})
.transpose()?;
if &actual != expected {
bail!("compiler environment input changed: {name}");
}
}
Ok(())
}
fn restore_result(
action: &RustcAction,
outputs: &RustcOutputs,
discovered: &DiscoveredInputs,
restore_outputs: bool,
mappings: &[PathMapping],
) -> Result<Option<CachedCompilation>> {
let responses = session::request_agent(&[AgentRequest::FindActionResult {
action: action.digest.clone(),
}])?;
let Some(response) = responses.into_iter().next() else {
bail!("cache agent did not return an action lookup response");
};
let result = match response {
AgentResponse::ActionResult { result } => match result {
Some(result) => result,
None => return Ok(None),
},
AgentResponse::Error { message } => bail!(message),
_ => bail!("cache agent returned an unexpected action lookup response"),
};
if result.version != 1 || result.action != action.digest {
bail!("cached rustc action result has an invalid identity");
}
let metadata_digest = result
.metadata
.ok_or_else(|| eyre::eyre!("cached rustc action result has no metadata"))?;
let output_root_digest = result
.output_root
.ok_or_else(|| eyre::eyre!("cached rustc action result has no output root"))?;
let roots = find_blobs(&[
action.digest.clone(),
metadata_digest.clone(),
output_root_digest.clone(),
])?;
let cached_action = read_verified_blob(&roots[0], &action.digest, "action descriptor")?;
if cached_action != action.bytes {
bail!("cached rustc action descriptor does not match the invocation");
}
let metadata: RustcMetadata =
read_canonical_blob(&roots[1], &metadata_digest, "rustc metadata")?;
if metadata.version != 1 || metadata.kind != "rustc" {
bail!("cached rustc metadata is unsupported");
}
let directory: CacheDirectory =
read_canonical_blob(&roots[2], &output_root_digest, "output directory")?;
let files = validated_outputs(directory, outputs)?;
let restored_output_files = files.len().try_into().unwrap_or(u64::MAX);
let restored_output_bytes = files.iter().fold(0_u64, |total, (node, _)| {
total.saturating_add(node.digest.size)
});
let mut digests = vec![metadata.stdout.clone(), metadata.stderr.clone()];
digests.extend(files.iter().map(|(node, _)| node.digest.clone()));
let blobs = find_blobs(&digests)?;
let stdout = denormalize_output_text(
&read_verified_blob(&blobs[0], &metadata.stdout, "stdout")?,
mappings,
);
let stderr = denormalize_output_text(
&read_verified_blob(&blobs[1], &metadata.stderr, "stderr")?,
mappings,
);
let materialization_started = Instant::now();
std::fs::create_dir_all(&outputs.directory)?;
let staging = tempfile::tempdir_in(&outputs.directory)?;
let mut staged = Vec::with_capacity(files.len());
let mut restore = RestoreStats {
output_files: restored_output_files,
output_bytes: restored_output_bytes,
..RestoreStats::default()
};
let mut cached_outputs = Vec::with_capacity(files.len());
for (index, ((node, destination), source)) in files.into_iter().zip(&blobs[2..]).enumerate() {
if destination == outputs.dep_info {
let bytes = denormalize_output_text(
&read_verified_blob(source, &node.digest, "dep-info")?,
mappings,
);
let digest = CacheDigest::blake3(&bytes);
cached_outputs.push(CachedOutput {
path: destination.clone(),
digest: digest.clone(),
executable: node.executable,
mode: node.mode,
});
if restore_outputs
&& std::fs::read(&destination).is_ok_and(|existing| existing == bytes)
{
restore.reused_output_files = restore.reused_output_files.saturating_add(1);
restore.reused_output_bytes = restore
.reused_output_bytes
.saturating_add(bytes.len().try_into().unwrap_or(u64::MAX));
continue;
}
let temporary = staging.path().join(format!("output-{index}"));
std::fs::write(&temporary, &bytes)?;
let temporary = tempfile::TempPath::try_from_path(temporary)?;
apply_file_mode(&temporary, node.mode, node.executable)?;
restore.copied_output_files = restore.copied_output_files.saturating_add(1);
restore.copied_output_bytes = restore
.copied_output_bytes
.saturating_add(bytes.len().try_into().unwrap_or(u64::MAX));
staged.push((temporary, destination));
continue;
}
cached_outputs.push(CachedOutput {
path: destination.clone(),
digest: node.digest.clone(),
executable: node.executable,
mode: node.mode,
});
if restore_outputs
&& output_already_in_place(&node, &destination, session::file_digest_cache())
{
restore.reused_output_files = restore.reused_output_files.saturating_add(1);
restore.reused_output_bytes =
restore.reused_output_bytes.saturating_add(node.digest.size);
continue;
}
let (temporary, materialization) =
stage_verified_cached_output(staging.path(), index, source, &node)?;
match materialization {
Materialization::Reflink => {
restore.reflinked_output_files = restore.reflinked_output_files.saturating_add(1);
restore.reflinked_output_bytes = restore
.reflinked_output_bytes
.saturating_add(node.digest.size);
}
Materialization::Copy => {
restore.copied_output_files = restore.copied_output_files.saturating_add(1);
restore.copied_output_bytes =
restore.copied_output_bytes.saturating_add(node.digest.size);
}
}
staged.push((temporary, destination));
}
let staged = StagedOutputs {
directory: staging,
files: staged,
};
verify_environment(&discovered.environment)?;
finalize_restored_outputs(staged, restore_outputs)?;
if restore_outputs {
record_output_digests(&cached_outputs);
}
restore.duration_ns = materialization_started
.elapsed()
.as_nanos()
.try_into()
.unwrap_or(u64::MAX);
Ok(Some(CachedCompilation {
stdout,
stderr,
outputs: cached_outputs,
restore,
}))
}
pub(crate) fn output_already_in_place(
node: &CacheFileNode,
destination: &Path,
digests: &dyn mbx_cache_core::FileDigestCache,
) -> bool {
let Ok(metadata) = std::fs::metadata(destination) else {
return false;
};
if !metadata.is_file()
|| metadata.len() != node.digest.size
|| !executable_mode_matches(&metadata, node.executable)
{
return false;
}
if let Some(identity) = FileIdentity::describe(destination, &metadata)
&& let Some(recorded) = digests
.find(FileDigestScope::Content, &[identity])
.pop()
.flatten()
{
return recorded == node.digest;
}
CacheDigest::blake3_file(destination).is_ok_and(|digest| digest == node.digest)
}
fn record_output_digests(outputs: &[CachedOutput]) {
let entries = outputs
.iter()
.filter_map(|output| {
let metadata = std::fs::metadata(&output.path).ok()?;
if metadata.len() != output.digest.size {
return None;
}
Some(RecordedFileDigest {
file: FileIdentity::describe(&output.path, &metadata)?,
digest: output.digest.clone(),
})
})
.collect::<Vec<_>>();
session::record_file_digests(FileDigestScope::Content, entries);
}
fn finalize_restored_outputs(staged: StagedOutputs, restore_outputs: bool) -> Result<()> {
if restore_outputs {
persist_outputs(staged)?;
}
Ok(())
}
fn verification_divergence(cached: &CachedCompilation, output: &Output) -> Option<String> {
if !output.status.success() {
return Some("the shadow compilation failed".into());
}
if cached.stdout != output.stdout {
return Some("standard output differs".into());
}
if cached.stderr != output.stderr {
return Some("standard error differs".into());
}
for expected in &cached.outputs {
let name = expected.path.display();
let Ok(metadata) = std::fs::metadata(&expected.path) else {
return Some(format!("{name} is missing"));
};
if file_mode(&metadata) != expected.mode {
return Some(format!("{name} has a different file mode"));
}
if !executable_mode_matches(&metadata, expected.executable) {
return Some(format!("{name} has a different executable bit"));
}
if !expected
.digest
.matches_file(&expected.path)
.unwrap_or(false)
{
return Some(format!("{name} has different contents"));
}
}
None
}
fn validated_outputs(
directory: CacheDirectory,
outputs: &RustcOutputs,
) -> Result<Vec<(CacheFileNode, PathBuf)>> {
if directory.version != 1 || !directory.directories.is_empty() || !directory.symlinks.is_empty()
{
bail!("cached rustc output directory has unsupported entries");
}
let mut expected = outputs
.files
.iter()
.chain(std::iter::once(&outputs.dep_info))
.map(|path| {
let name = path
.file_name()
.and_then(|name| name.to_str())
.ok_or_else(|| eyre::eyre!("expected rustc output name is not UTF-8"))?;
if path.parent() != Some(outputs.directory.as_path()) {
bail!("expected rustc output escapes its output directory");
}
Ok((
name.to_string(),
(path.clone(), outputs.is_executable(path)),
))
})
.collect::<Result<BTreeMap<_, _>>>()?;
if directory.files.len() != expected.len() {
bail!("cached rustc output set does not match the invocation");
}
let mut files = Vec::with_capacity(directory.files.len());
for node in directory.files {
let (destination, executable) = expected
.remove(&node.name)
.ok_or_else(|| eyre::eyre!("cached rustc output is unexpected: {}", node.name))?;
validate_file_mode(&node, executable)?;
files.push((node, destination));
}
if !expected.is_empty() {
bail!("cached rustc output set is incomplete");
}
Ok(files)
}
fn compiler_identity(rustc: &OsStr) -> Result<CompilerIdentity> {
static IDENTITY: OnceLock<(OsString, CompilerIdentity)> = OnceLock::new();
if let Some((known, identity)) = IDENTITY.get()
&& known.as_os_str() == rustc
{
return Ok(identity.clone());
}
let identity = query_compiler_identity(rustc)?;
let _ = IDENTITY.set((rustc.to_os_string(), identity.clone()));
Ok(identity)
}
fn query_compiler_identity(rustc: &OsStr) -> Result<CompilerIdentity> {
let executable = resolve_executable(rustc)?;
let environment = ["RUSTUP_HOME", "RUSTUP_TOOLCHAIN"]
.into_iter()
.map(|name| (name.into(), std::env::var(name).ok()))
.collect::<BTreeMap<_, _>>();
let responses = session::request_agent(&[AgentRequest::FindExecutableIdentity {
executable: executable.clone(),
environment: environment.clone(),
}])?;
let Some(AgentResponse::ExecutableIdentity { stdout }) = responses.into_iter().next() else {
bail!("cache agent did not return the rustc identity");
};
let stdout = if let Some(stdout) = stdout {
stdout
} else {
let mut command = Command::new(&executable);
command.arg("-vV");
for (name, value) in &environment {
if let Some(value) = value {
command.env(name, value);
} else {
command.env_remove(name);
}
}
let output = command
.output()
.wrap_err("failed to query the rustc identity")?;
if !output.status.success() {
bail!(
"rustc identity command failed: {}",
String::from_utf8_lossy(&output.stderr)
);
}
let responses = session::request_agent(&[AgentRequest::StoreExecutableIdentity {
executable,
environment,
stdout: output.stdout,
}])?;
let Some(AgentResponse::ExecutableIdentity {
stdout: Some(stdout),
}) = responses.into_iter().next()
else {
bail!("cache agent did not store the rustc identity");
};
stdout
};
let verbose = std::str::from_utf8(&stdout).wrap_err("rustc identity is not UTF-8")?;
let release = identity_field(verbose, "release")?;
let host = identity_field(verbose, "host")?;
let rustc_version = verbose
.lines()
.filter(|line| {
line.starts_with("rustc ")
|| line.starts_with("commit-hash:")
|| line.starts_with("commit-date:")
|| line.starts_with("LLVM version:")
})
.collect::<Vec<_>>()
.join("; ");
if rustc_version.is_empty() {
bail!("rustc identity is missing its version");
}
let toolchain = std::env::var("RUSTUP_TOOLCHAIN")
.ok()
.filter(|value| !value.is_empty())
.unwrap_or_else(|| release.to_string());
Ok(CompilerIdentity {
toolchain,
rustc_version,
host: host.to_string(),
})
}
fn identity_field<'a>(verbose: &'a str, field: &str) -> Result<&'a str> {
verbose
.lines()
.find_map(|line| line.strip_prefix(&format!("{field}: ")))
.filter(|value| !value.is_empty())
.ok_or_else(|| eyre::eyre!("rustc identity is missing {field}"))
}
const PORTABLE_ENVIRONMENT: &[&str] = &["OUT_DIR"];
struct Portable {
mappings: Vec<PathMapping>,
arguments: Vec<OsString>,
names: BTreeSet<String>,
values: Vec<String>,
}
impl Portable {
fn detect(working_dir: &Path, target_output: Option<&Path>, target: Option<&str>) -> Self {
let mut portable = Self {
mappings: PathMapping::ordered(&path_mappings(working_dir, target_output, target)),
arguments: Vec::new(),
names: BTreeSet::new(),
values: Vec::new(),
};
if !session::share_out_dir_requested() {
return portable;
}
for name in PORTABLE_ENVIRONMENT {
let Some(value) = std::env::var(name)
.ok()
.filter(|value| Path::new(value).is_absolute())
else {
continue;
};
let Ok(placeholder) =
normalize_mapped_path(Path::new(&value), working_dir, &portable.mappings)
else {
continue;
};
let mut flag = OsString::from("--remap-path-prefix=");
flag.push(&value);
flag.push("=");
flag.push(&placeholder);
portable.arguments.push(flag);
portable.names.insert((*name).to_string());
portable.values.push(value);
}
portable
}
fn applied_to(&self, arguments: &[OsString]) -> Vec<OsString> {
let mut applied = arguments.to_vec();
applied.extend(self.arguments.iter().cloned());
applied
}
fn contents_are_clean(&self, contents: &[u8]) -> bool {
!self.values.is_empty() && !self.values.iter().any(|value| carries(contents, value))
}
}
fn carries(contents: &[u8], value: &str) -> bool {
if contains(contents, value.as_bytes()) {
return true;
}
value.contains('\\') && contains(contents, value.replace('\\', "/").as_bytes())
}
fn contains(haystack: &[u8], needle: &[u8]) -> bool {
!needle.is_empty() && memchr::memmem::find(haystack, needle).is_some()
}
fn path_mappings(
working_dir: &Path,
target_output: Option<&Path>,
target: Option<&str>,
) -> Vec<PathMapping> {
path_mappings_with_env(working_dir, target_output, target, |name| {
std::env::var_os(name)
})
}
fn path_mappings_with_env(
working_dir: &Path,
target_output: Option<&Path>,
target: Option<&str>,
environment: impl Fn(&str) -> Option<OsString>,
) -> Vec<PathMapping> {
let mut mappings = Vec::new();
let mut roots = BTreeSet::new();
let home_roots = ["HOME", "USERPROFILE"]
.into_iter()
.filter_map(|name| environment(name).map(PathBuf::from))
.filter(|root| root.is_absolute())
.collect::<Vec<_>>();
let configured_target = environment(session::TARGET_DIR_ENV)
.map(PathBuf::from)
.filter(|root| root.is_absolute());
if let Some(root) = configured_target.or_else(|| {
target_output
.filter(|root| root.is_absolute())
.map(|output| standalone_target_root(output, target))
}) {
add_mapping(&mut mappings, &mut roots, root, "target");
}
for (name, placeholder) in [
(session::WORKSPACE_ROOT_ENV, "workspace"),
("CARGO_HOME", "cargo_home"),
("RUSTUP_HOME", "rustup_home"),
] {
if let Some(root) = environment(name).map(PathBuf::from)
&& root.is_absolute()
{
add_mapping(&mut mappings, &mut roots, root, placeholder);
}
}
if let Some(home) = home_roots.first() {
for (directory, placeholder) in [(".cargo", "cargo_home"), (".rustup", "rustup_home")] {
if !mappings
.iter()
.any(|mapping| mapping.placeholder == placeholder)
{
add_mapping(&mut mappings, &mut roots, home.join(directory), placeholder);
}
}
}
if !mappings
.iter()
.any(|mapping| mapping.placeholder == "workspace")
&& !roots.iter().any(|root| working_dir.starts_with(root))
{
add_mapping(
&mut mappings,
&mut roots,
workspace_root(working_dir),
"workspace",
);
}
for root in home_roots {
add_mapping(&mut mappings, &mut roots, root, "home");
}
mappings
}
fn standalone_target_root(output: &Path, target: Option<&str>) -> PathBuf {
if output.file_name() == Some(OsStr::new("deps"))
&& let Some(profile_root) = output.parent().and_then(Path::parent)
{
let target_component = target.and_then(|target| Path::new(target).file_stem());
if target_component.is_some_and(|target| profile_root.file_name() == Some(target))
&& let Some(root) = profile_root.parent()
{
return root.to_path_buf();
}
return profile_root.to_path_buf();
}
output.to_path_buf()
}
fn add_mapping(
mappings: &mut Vec<PathMapping>,
roots: &mut BTreeSet<PathBuf>,
root: PathBuf,
placeholder: &str,
) {
if roots.insert(root.clone())
&& !mappings
.iter()
.any(|mapping| mapping.placeholder == placeholder)
{
mappings.push(PathMapping::new(root, placeholder));
}
}
fn replay_output(output: &Output) -> Result<()> {
replay_bytes(&output.stdout, &output.stderr)
}
fn publish_result<'a>(
candidates: &'a ActionCandidates,
portable: &Portable,
outputs: &RustcOutputs,
output: &Output,
mappings: &[PathMapping],
) -> Result<&'a RustcAction> {
if outputs.files.is_empty() {
bail!("rustc produced no cacheable outputs");
}
let staging = staging_directory()?;
let mut blobs = Vec::new();
let stdout = staged_bytes(
staging.path(),
"stdout",
&normalize_output_text(&output.stdout, mappings),
)?;
let stderr = staged_bytes(
staging.path(),
"stderr",
&normalize_output_text(&output.stderr, mappings),
)?;
blobs.extend([stdout.clone(), stderr.clone()]);
let output_paths = outputs
.files
.iter()
.chain(std::iter::once(&outputs.dep_info));
let mut files = Vec::with_capacity(outputs.files.len() + 1);
let mut hashed_outputs = Vec::with_capacity(outputs.files.len());
let mut portable_outputs_are_clean = candidates.portable.is_some();
for path in output_paths {
let metadata = std::fs::metadata(path)
.wrap_err_with(|| format!("failed to inspect rustc output {}", path.display()))?;
if !metadata.is_file() {
bail!("rustc output is not a regular file: {}", path.display());
}
let digest = if path == &outputs.dep_info {
let normalized = normalize_output_text(&std::fs::read(path)?, mappings);
let staged = staged_bytes(staging.path(), "dep-info", &normalized)?;
blobs.push(staged.clone());
staged.0
} else {
let digest = if candidates.portable.is_some() {
let contents = std::fs::read(path)
.wrap_err_with(|| format!("failed to read rustc output {}", path.display()))?;
portable_outputs_are_clean &= portable.contents_are_clean(&contents);
CacheDigest::blake3(&contents)
} else {
CacheDigest::blake3_file(path)?
};
blobs.push((digest.clone(), path.clone()));
if metadata.len() == digest.size
&& let Some(file) = FileIdentity::describe(path, &metadata)
{
hashed_outputs.push(RecordedFileDigest {
file,
digest: digest.clone(),
});
}
digest
};
files.push(CacheFileNode {
digest,
executable: outputs.is_executable(path),
mode: file_mode(&metadata),
name: path
.file_name()
.and_then(|name| name.to_str())
.ok_or_else(|| eyre::eyre!("rustc output name is not UTF-8"))?
.to_string(),
});
}
files.sort_by(|left, right| left.name.cmp(&right.name));
session::record_file_digests(FileDigestScope::Content, hashed_outputs);
let action = candidates.publishable(portable_outputs_are_clean);
blobs.push(staged_bytes(staging.path(), "action.json", &action.bytes)?);
let metadata = canonical_json(&RustcMetadata {
version: 1,
kind: "rustc".into(),
stdout: stdout.0,
stderr: stderr.0,
})?;
let metadata = staged_bytes(staging.path(), "metadata.json", &metadata)?;
blobs.push(metadata.clone());
let directory = canonical_json(&CacheDirectory {
directories: Vec::new(),
files,
symlinks: Vec::new(),
version: 1,
})?;
let directory = staged_bytes(staging.path(), "directory.json", &directory)?;
blobs.push(directory.clone());
let mut requests = Vec::new();
let mut published = BTreeSet::new();
for (digest, source) in blobs {
if published.insert(digest.clone()) {
requests.push(AgentRequest::StoreBlob { digest, source });
}
}
requests.push(AgentRequest::StoreActionResult {
result: RemoteActionResult {
action: action.digest.clone(),
metadata: Some(metadata.0),
output_root: Some(directory.0),
version: 1,
},
});
for response in session::request_agent(&requests)? {
match response {
AgentResponse::Stored { .. } | AgentResponse::ActionStored { .. } => {}
AgentResponse::Error { message } => bail!(message),
_ => bail!("cache agent returned an unexpected publish response"),
}
}
Ok(action)
}
fn staged_bytes(directory: &Path, name: &str, bytes: &[u8]) -> Result<(CacheDigest, PathBuf)> {
let path = directory.join(name);
std::fs::write(&path, bytes)?;
Ok((CacheDigest::blake3(bytes), path))
}
#[cfg(test)]
#[path = "rustc_tests.rs"]
mod tests;