1use crate::uploads::{ConnectionUploads, UploadQueue, UploadSink};
2use crate::{
3 ActionPrediction, ActionPromiseCompletion, ActionPromiseState, CacheDigest, CacheDirectory,
4 LocalActionCache, LocalCas, MAX_STAGED_BLOB_PACK_BYTES, ManifestPutOutcome, RemoteActionResult,
5 RemoteCacheClient, RemoteCacheMode, RustcMetadata, TaskActionManifest, blob_pack_chunk,
6 canonical_json,
7};
8use eyre::{Context, Result, bail};
9use futures_util::{FutureExt, StreamExt, future::BoxFuture, stream};
10use log::{debug, warn};
11use serde::{Deserialize, Serialize};
12use std::collections::{BTreeMap, BTreeSet};
13use std::fs;
14use std::path::{Path, PathBuf};
15use std::sync::atomic::{AtomicU64, Ordering};
16use std::sync::{Arc, Mutex, Weak};
17use std::time::{Duration, Instant, SystemTime};
18
19const ACTION_PROMISE_WAIT: Duration = Duration::from_secs(60 * 60);
23use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
24
25mod file_digest;
26mod manifest;
27mod prefetch;
28mod stats;
29mod wire;
30
31#[cfg(test)]
32pub(crate) use prefetch::select_prefetch_actions;
33
34pub use file_digest::{
35 FileDigestCache, FileDigestResolution, FileDigestScope, FileIdentity, FileObjectIdentity,
36 FileSnapshot, NoFileDigestCache, RecordedFileDigest, ValidatedFileDigest, digest_file,
37 digest_file_validated,
38};
39pub use manifest::{
40 is_task_identity, task_manifest_actions,
41 update_task_predictions_in_order as merge_task_action_predictions,
42};
43use manifest::{
44 merge_task_manifests, task_manifest_dir, update_task_predictions, validate_task_identity,
45 validate_task_manifest,
46};
47pub use stats::{AgentStats, CompilerStats};
48use wire::MAX_REQUEST_BYTES;
49pub use wire::{
50 AGENT_PROTOCOL_VERSION, ActionDiagnostic, AgentEvent, AgentEventObserver, AgentRequest,
51 AgentResponse, PinnedFile, PinnedState, RestoreStats, WrapperSpan, WrapperTiming,
52};
53
54const MAX_EXECUTABLE_IDENTITIES: usize = 64;
55const MAX_EXECUTABLE_IDENTITY_SIZE: usize = 64 * 1024;
56const MAX_EXECUTABLE_IDENTITY_BYTES: usize = 256 * 1024;
57const TASK_ACTION_MANIFEST_VERSION: u8 = 1;
58const MAX_TASK_ACTION_PREDICTIONS: usize = 16 * 1024;
59const MAX_TASK_MANIFEST_BYTES: usize = 15 * 1024 * 1024;
62const TASK_ACTION_LRU_REFRESH_MARGIN: usize = MAX_TASK_ACTION_PREDICTIONS / 16;
63const MAX_DIAGNOSTIC_BYTES: usize = 4 * 1024;
65const MAX_DIAGNOSTICS: usize = 128;
71const ACTION_DIAGNOSTIC_PREFIX: &str = "@mbx-action-diagnostic\t";
72
73#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
74enum DiagnosticLevel {
75 Warning,
76 Error,
77}
78
79impl DiagnosticLevel {
80 fn label(self) -> &'static str {
81 match self {
82 Self::Warning => "warning",
83 Self::Error => "error",
84 }
85 }
86}
87
88const MAX_FILE_DIGEST_BATCH: usize = 16 * 1024;
90const MAX_FILE_DIGEST_ENTRIES: usize = 1024 * 1024;
95const MAX_CONCURRENT_FILE_DIGESTS: usize = 8;
98const MAX_CONCURRENT_BLOB_VERIFICATIONS: usize = 16;
102const MAX_REMOTE_TRANSFERS: usize = 64;
103const MAX_PREFETCH_TRANSFERS: usize = 48;
104const MAX_CONCURRENT_BLOB_PACKS: usize = 8;
110const MAX_PREFETCH_ACTIONS: usize = 1024;
121const MAX_PREFETCH_ACTION_WAVE: usize = 32;
126const MAX_PREFETCH_BATCH_LOOKUPS: usize = 1;
132const PREFETCH_ACTION_BATCH_DELAY: Duration = Duration::from_millis(5);
133const MAX_PREFETCH_DIRECTORY_OBJECTS: usize = 100_000;
134const MAX_PREFETCH_OBJECTS_PER_WAVE: usize = 100_000;
135const DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES: u64 = 5 * 1024 * 1024 * 1024;
136
137pub type TaskFallbacks = Arc<dyn Fn() -> Vec<String> + Send + Sync>;
140type FileDigestFlightKey = (FileDigestScope, FileIdentity);
141type FileDigestFlights = BTreeMap<FileDigestFlightKey, Weak<FileDigestFlight>>;
142
143struct FileDigestFlight {
144 lock: tokio::sync::Mutex<()>,
145 resolution: Mutex<Option<FileDigestResolution>>,
146}
147
148const NEWEST_MANIFEST_CANDIDATES: usize = 8;
160
161pub struct AgentRemoteCache {
163 pub client: RemoteCacheClient,
165 pub mode: RemoteCacheMode,
167 pub staging_dir: PathBuf,
169}
170
171#[derive(Default)]
172struct AtomicAgentStats {
173 wrapper_phases_ns: Mutex<BTreeMap<String, u64>>,
174 lookups: AtomicU64,
175 unconsulted: AtomicU64,
176 incremental_compilations: AtomicU64,
177 hits: AtomicU64,
178 stores: AtomicU64,
179 stored_bytes: AtomicU64,
180 verifications: AtomicU64,
181 divergences: AtomicU64,
182 downloaded_bytes: AtomicU64,
183 uploaded_bytes: AtomicU64,
184 background_uploads: AtomicU64,
185 background_upload_failures: AtomicU64,
186 remote_blob_pack_uploads: AtomicU64,
187 remote_blob_pack_upload_blobs: AtomicU64,
188 upload_drain_duration_ns: AtomicU64,
189 prefetched_actions: AtomicU64,
190 predictions_loaded: AtomicU64,
191 remote_failures: AtomicU64,
192 remote_manifest_lookups: AtomicU64,
193 remote_manifest_lookup_duration_ns: AtomicU64,
194 remote_action_lookups: AtomicU64,
195 remote_action_lookup_duration_ns: AtomicU64,
196 remote_blob_requests: AtomicU64,
197 remote_blob_pack_requests: AtomicU64,
198 remote_blob_pack_blobs: AtomicU64,
199 remote_blob_transfer_duration_ns: AtomicU64,
200 local_cas_write_duration_ns: AtomicU64,
201 prefetch_runs: AtomicU64,
202 prefetch_duration_ns: AtomicU64,
203 materialization_duration_ns: AtomicU64,
204 bypasses: Mutex<BTreeMap<String, u64>>,
205 avoided_compiler_duration_ns: AtomicU64,
206 compiler: Mutex<BTreeMap<String, CompilerStats>>,
207 unit_outcomes: Mutex<BTreeMap<String, BTreeSet<String>>>,
208 slow_compilations: Mutex<BTreeMap<String, u64>>,
209 restored_output_files: AtomicU64,
210 restored_output_bytes: AtomicU64,
211 reflinked_output_files: AtomicU64,
212 reflinked_output_bytes: AtomicU64,
213 hardlinked_output_files: AtomicU64,
214 hardlinked_output_bytes: AtomicU64,
215 copied_output_files: AtomicU64,
216 copied_output_bytes: AtomicU64,
217 reused_output_files: AtomicU64,
218 reused_output_bytes: AtomicU64,
219}
220
221struct AtomicDurationTimer<'a> {
222 started: Instant,
223 target: &'a AtomicU64,
224}
225
226impl<'a> AtomicDurationTimer<'a> {
227 fn start(target: &'a AtomicU64) -> Self {
228 Self {
229 started: Instant::now(),
230 target,
231 }
232 }
233}
234
235impl Drop for AtomicDurationTimer<'_> {
236 fn drop(&mut self) {
237 atomic_saturating_add(self.target, duration_ns(self.started));
238 }
239}
240
241fn duration_ns(started: Instant) -> u64 {
242 started.elapsed().as_nanos().try_into().unwrap_or(u64::MAX)
243}
244
245fn validate_crate_name(crate_name: Option<&str>) -> Result<()> {
246 if let Some(crate_name) = crate_name
247 && (crate_name.len() > 256 || crate_name.contains(['\0', '\n', '\r']))
248 {
249 bail!("invalid compiler crate name");
250 }
251 Ok(())
252}
253
254#[derive(serde::Deserialize)]
255struct ActionDiagnosticEnvelope {
256 outcome: String,
257 crate_name: Option<String>,
258 diagnostic: ActionDiagnostic,
259}
260
261pub(crate) fn lookup_outcome(outcome: &str) -> (&str, bool) {
268 match outcome {
269 "incremental-miss" => ("miss", true),
270 "incremental-unconsulted" => ("unconsulted", true),
271 other => (other, false),
272 }
273}
274
275fn parse_action_diagnostic(
276 message: &str,
277) -> Option<Result<(String, Option<String>, ActionDiagnostic)>> {
278 let payload = message.strip_prefix(ACTION_DIAGNOSTIC_PREFIX)?;
279 Some((|| {
280 let envelope: ActionDiagnosticEnvelope = serde_json::from_str(payload)?;
281 let (outcome, _) = lookup_outcome(&envelope.outcome);
285 if !matches!(outcome, "hit" | "miss" | "unconsulted") {
286 bail!("invalid action diagnostic outcome");
287 }
288 validate_crate_name(envelope.crate_name.as_deref())?;
289 Ok((
290 outcome.to_string(),
291 envelope.crate_name,
292 envelope.diagnostic,
293 ))
294 })())
295}
296
297fn atomic_saturating_add(target: &AtomicU64, value: u64) {
298 let mut current = target.load(Ordering::Relaxed);
301 while let Err(actual) = target.compare_exchange_weak(
302 current,
303 current.saturating_add(value),
304 Ordering::Relaxed,
305 Ordering::Relaxed,
306 ) {
307 current = actual;
308 }
309}
310
311fn queue_prefetch_digest(
312 verified: &BTreeMap<CacheDigest, PathBuf>,
313 pending: &mut BTreeMap<CacheDigest, ()>,
314 digest: CacheDigest,
315) {
316 if verified.contains_key(&digest) || pending.contains_key(&digest) {
317 return;
318 }
319 pending.insert(digest, ());
320}
321
322fn queue_prefetch_directory(
323 seen: &BTreeMap<CacheDigest, ()>,
324 pending: &mut BTreeMap<CacheDigest, ()>,
325 digest: CacheDigest,
326 limit: usize,
327) -> bool {
328 if seen.contains_key(&digest) || pending.contains_key(&digest) {
329 return true;
330 }
331 if seen.len().saturating_add(pending.len()) >= limit {
332 return false;
333 }
334 pending.insert(digest, ());
335 true
336}
337
338pub type TaskLoader = dyn for<'a> Fn(&'a CacheAgent, &'a str) -> BoxFuture<'a, ()> + Send + Sync;
340
341#[derive(Clone)]
347pub struct CacheAgent {
348 cas: LocalCas,
349 actions: LocalActionCache,
350 verified_blobs: Arc<Mutex<BTreeMap<CacheDigest, VerifiedBlob>>>,
351 version: Arc<str>,
352 write_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
353 action_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
354 stats: Arc<AtomicAgentStats>,
355 observer: Option<Arc<dyn AgentEventObserver>>,
356 observer_emission: Arc<Mutex<()>>,
357 task_loader: Option<Arc<TaskLoader>>,
360 executable_identities: Arc<Mutex<BTreeMap<ExecutableIdentityKey, Vec<u8>>>>,
361 identity_dir: Arc<PathBuf>,
363 manifest_dir: Arc<PathBuf>,
364 task_actions: Arc<Mutex<BTreeMap<String, TaskActionState>>>,
365 task_fallbacks: Arc<Mutex<BTreeMap<String, TaskFallbacks>>>,
367 next_task_run: Arc<AtomicU64>,
368 manifest_write_lock: Arc<Mutex<()>>,
369 remote: Option<Arc<RemoteCacheClient>>,
370 remote_mode: RemoteCacheMode,
371 remote_staging_dir: Arc<PathBuf>,
372 remote_download_limit: u64,
373 remote_download_bytes: Arc<AtomicU64>,
374 pending_remote_actions: Arc<Mutex<BTreeMap<CacheDigest, RemoteActionResult>>>,
375 remote_transfers: Arc<tokio::sync::Semaphore>,
376 prefetch_transfers: Arc<tokio::sync::Semaphore>,
377 prefetch_tasks: Arc<Mutex<Vec<tokio::task::JoinHandle<()>>>>,
378 diagnostics: Arc<Mutex<BTreeSet<(DiagnosticLevel, String)>>>,
383 file_digests: Arc<Mutex<BTreeMap<(FileDigestScope, PathBuf), RecordedFileDigest>>>,
386 file_digest_locks: Arc<Mutex<FileDigestFlights>>,
388 file_digest_permits: Arc<tokio::sync::Semaphore>,
389 blob_verification_permits: Arc<tokio::sync::Semaphore>,
390 #[cfg(test)]
391 file_digest_reads: Arc<AtomicU64>,
392 uploads: Option<UploadQueue>,
394}
395
396fn validate_file_digest_record(entry: &RecordedFileDigest) -> Result<()> {
399 if !entry.file.path.is_absolute() {
400 bail!("file-digest records need absolute paths");
401 }
402 entry.digest.validate()?;
403 if entry.file.len != entry.digest.size {
404 bail!("file-digest record length does not match its digest");
405 }
406 Ok(())
407}
408
409struct AgentUploadSink {
411 stats: Arc<AtomicAgentStats>,
412}
413
414impl UploadSink for AgentUploadSink {
415 fn record_blob_uploaded(&self, bytes: u64) {
416 self.stats
417 .background_uploads
418 .fetch_add(1, Ordering::Relaxed);
419 self.stats
420 .uploaded_bytes
421 .fetch_add(bytes, Ordering::Relaxed);
422 }
423
424 fn record_action_uploaded(&self) {
425 self.stats
426 .background_uploads
427 .fetch_add(1, Ordering::Relaxed);
428 }
429
430 fn record_blob_pack_uploaded(&self, blobs: u64) {
431 self.stats
432 .remote_blob_pack_uploads
433 .fetch_add(1, Ordering::Relaxed);
434 self.stats
435 .remote_blob_pack_upload_blobs
436 .fetch_add(blobs, Ordering::Relaxed);
437 }
438
439 fn record_upload_failure(&self) {
440 self.stats
441 .background_upload_failures
442 .fetch_add(1, Ordering::Relaxed);
443 self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
444 }
445}
446
447#[derive(Debug, Clone)]
453struct VerifiedBlob {
454 path: PathBuf,
455 len: u64,
456 modified: SystemTime,
457 object: Option<(u64, u64)>,
460}
461
462impl VerifiedBlob {
463 fn describe(path: &Path) -> Option<Self> {
466 let metadata = std::fs::metadata(path).ok()?;
467 Some(Self {
468 path: path.to_path_buf(),
469 len: metadata.len(),
470 modified: metadata.modified().ok()?,
471 object: file_object(&metadata),
472 })
473 }
474
475 fn is_unchanged(&self) -> bool {
477 let Ok(metadata) = std::fs::metadata(&self.path) else {
478 return false;
479 };
480 metadata.len() == self.len
481 && metadata.modified().is_ok_and(|now| now == self.modified)
482 && file_object(&metadata) == self.object
483 }
484}
485
486#[cfg(unix)]
488fn file_object(metadata: &std::fs::Metadata) -> Option<(u64, u64)> {
489 use std::os::unix::fs::MetadataExt as _;
490 Some((metadata.dev(), metadata.ino()))
491}
492
493#[cfg(not(unix))]
494fn file_object(_metadata: &std::fs::Metadata) -> Option<(u64, u64)> {
495 None
496}
497
498fn find_stored_blob(cas: &LocalCas, digest: &CacheDigest) -> Option<PathBuf> {
500 cas.find(digest).unwrap_or_else(|error| {
507 warn!(
508 "local CAS blob {} is not restorable and will be republished: {error}",
509 digest.hash
510 );
511 None
512 })
513}
514
515#[derive(Debug, Clone, Default)]
516struct TaskActionState {
517 manifest: String,
518 baseline_loaded: bool,
519 predictions: BTreeMap<CacheDigest, ActionPrediction>,
520 pending_predictions: BTreeMap<CacheDigest, ActionPrediction>,
521 changed_predictions: BTreeMap<CacheDigest, Option<ActionPrediction>>,
528 remote_predictions: Option<BTreeMap<CacheDigest, ActionPrediction>>,
533 remote_manifest: Option<TaskActionManifest>,
537 remote_publication_needed: bool,
541 prefetched_adapters: BTreeSet<String>,
542 remote_etag: Option<String>,
543}
544
545fn activate_prediction_adapter(
547 state: &mut TaskActionState,
548 invocation: &CacheDigest,
549) -> (Option<ActionPrediction>, Option<Vec<ActionPrediction>>) {
550 let prediction = state.predictions.get(invocation).cloned();
551 let prefetch = prediction.as_ref().and_then(|prediction| {
552 if !state.prefetched_adapters.insert(prediction.adapter.clone()) {
553 return None;
554 }
555 Some(
556 state
557 .predictions
558 .values()
559 .filter(|candidate| candidate.adapter == prediction.adapter)
560 .cloned()
561 .collect(),
562 )
563 });
564 (prediction, prefetch)
565}
566
567fn required_local_predictions(
570 predictions: &[ActionPrediction],
571 updates: &BTreeMap<CacheDigest, ActionPrediction>,
572 concurrent_updates: &BTreeMap<CacheDigest, ActionPrediction>,
573) -> BTreeSet<CacheDigest> {
574 let mut required: BTreeSet<_> = updates.keys().cloned().collect();
575 for prediction in predictions.iter().rev() {
576 if required.len() == MAX_TASK_ACTION_PREDICTIONS {
577 break;
578 }
579 if concurrent_updates.contains_key(&prediction.invocation) {
580 required.insert(prediction.invocation.clone());
581 }
582 }
583 required
584}
585
586fn concurrent_task_predictions(
589 predictions: &[ActionPrediction],
590 state: &TaskActionState,
591) -> BTreeMap<CacheDigest, ActionPrediction> {
592 predictions
593 .iter()
594 .filter(|prediction| {
595 let baseline = match state.changed_predictions.get(&prediction.invocation) {
596 Some(baseline) => baseline.as_ref(),
597 None => state.predictions.get(&prediction.invocation),
598 };
599 baseline != Some(*prediction)
600 })
601 .map(|prediction| (prediction.invocation.clone(), prediction.clone()))
602 .collect()
603}
604
605fn effective_task_updates(
610 predictions: &[ActionPrediction],
611 state: &TaskActionState,
612) -> BTreeMap<CacheDigest, ActionPrediction> {
613 let current: BTreeMap<_, _> = predictions
614 .iter()
615 .map(|prediction| (prediction.invocation.clone(), prediction))
616 .collect();
617 state
618 .pending_predictions
619 .iter()
620 .map(|(invocation, update)| {
621 let unchanged_hit = !state.changed_predictions.contains_key(invocation);
622 let effective = if unchanged_hit {
623 current.get(invocation).copied().unwrap_or(update).clone()
624 } else {
625 update.clone()
626 };
627 (invocation.clone(), effective)
628 })
629 .collect()
630}
631
632fn should_refresh_task_prediction_order(prediction_count: usize, has_recorded: bool) -> bool {
633 has_recorded && prediction_count >= MAX_TASK_ACTION_PREDICTIONS - TASK_ACTION_LRU_REFRESH_MARGIN
634}
635
636fn replaced_task_predictions(
639 remote: &BTreeMap<CacheDigest, ActionPrediction>,
640 candidates: &BTreeSet<CacheDigest>,
641) -> Vec<ActionPrediction> {
642 candidates
643 .iter()
644 .filter_map(|invocation| remote.get(invocation).cloned())
645 .collect()
646}
647
648fn remote_publication_precedence(
654 state: &TaskActionState,
655 completed: &[ActionPrediction],
656 concurrent_updates: &BTreeMap<CacheDigest, ActionPrediction>,
657 proven_uploads: &BTreeSet<CacheDigest>,
658 published: &BTreeMap<CacheDigest, &ActionPrediction>,
659 withheld: &BTreeSet<CacheDigest>,
660) -> (BTreeSet<CacheDigest>, BTreeSet<CacheDigest>) {
661 let publication_base_invocations: BTreeSet<_> = state
662 .remote_manifest
663 .iter()
664 .flat_map(|manifest| manifest.predictions.iter())
665 .map(|prediction| prediction.invocation.clone())
666 .collect();
667 let unchanged_hits: BTreeSet<_> = completed
668 .iter()
669 .filter(|prediction| {
670 !state
671 .changed_predictions
672 .contains_key(&prediction.invocation)
673 })
674 .map(|prediction| prediction.invocation.clone())
675 .collect();
676 let required = completed
677 .iter()
678 .filter(|prediction| !unchanged_hits.contains(&prediction.invocation))
679 .chain(concurrent_updates.values().filter(|prediction| {
680 state.remote_predictions.is_some() || proven_uploads.contains(&prediction.action)
681 }))
682 .filter(|prediction| {
683 !withheld.contains(&prediction.invocation)
684 && published.get(&prediction.invocation).copied() == Some(*prediction)
685 })
686 .map(|prediction| prediction.invocation.clone())
687 .collect();
688 let fallbacks = withheld
689 .iter()
690 .filter(|invocation| published.contains_key(*invocation))
691 .cloned()
692 .chain(
693 unchanged_hits
694 .into_iter()
695 .filter(|invocation| published.contains_key(invocation)),
696 )
697 .chain(
698 concurrent_updates
699 .iter()
700 .filter(|(invocation, prediction)| {
701 state.remote_predictions.is_none()
702 && !proven_uploads.contains(&prediction.action)
703 && published.contains_key(*invocation)
704 })
705 .map(|(invocation, _)| invocation.clone()),
706 )
707 .chain(
708 published
709 .keys()
710 .filter(|invocation| {
711 state.remote_predictions.is_none()
712 && state.remote_manifest.is_some()
713 && !publication_base_invocations.contains(*invocation)
714 })
715 .cloned(),
716 )
717 .collect();
718 (required, fallbacks)
719}
720
721fn unpublished_task_actions(
724 mut unproven: BTreeSet<CacheDigest>,
725 held: &BTreeSet<CacheDigest>,
726 failed: BTreeSet<CacheDigest>,
727) -> BTreeSet<CacheDigest> {
728 unproven.retain(|action| !held.contains(action));
729 unproven.extend(failed);
730 unproven
731}
732
733fn unproven_task_predictions<'a>(
735 predictions: impl IntoIterator<Item = &'a ActionPrediction>,
736 remote: &BTreeMap<CacheDigest, ActionPrediction>,
737) -> BTreeMap<CacheDigest, CacheDigest> {
738 predictions
739 .into_iter()
740 .filter(|prediction| remote.get(&prediction.invocation) != Some(*prediction))
741 .map(|prediction| (prediction.invocation.clone(), prediction.action.clone()))
742 .collect()
743}
744
745fn withhold_unpublished_updates(
748 manifest: &mut TaskActionManifest,
749 candidates: &BTreeSet<CacheDigest>,
750 replaced: &[ActionPrediction],
751 unpublished: &BTreeSet<CacheDigest>,
752) -> BTreeSet<CacheDigest> {
753 let withheld: BTreeSet<_> = manifest
754 .predictions
755 .iter()
756 .filter(|prediction| {
757 candidates.contains(&prediction.invocation) && unpublished.contains(&prediction.action)
758 })
759 .map(|prediction| prediction.invocation.clone())
760 .collect();
761 manifest
762 .predictions
763 .retain(|prediction| !withheld.contains(&prediction.invocation));
764 let mut retained: BTreeSet<_> = manifest
765 .predictions
766 .iter()
767 .map(|prediction| prediction.invocation.clone())
768 .collect();
769 for previous in replaced.iter().rev() {
770 if withheld.contains(&previous.invocation) && retained.insert(previous.invocation.clone()) {
771 manifest.predictions.push(previous.clone());
775 }
776 }
777 withheld
778}
779
780struct PrefetchedAction {
781 adapter: String,
782 result: RemoteActionResult,
783}
784
785struct RemoteDownloadReservation {
786 counter: Arc<AtomicU64>,
787 reserved: u64,
788 committed: bool,
789}
790
791impl RemoteDownloadReservation {
792 fn bytes(&self) -> u64 {
793 self.reserved
794 }
795
796 fn commit(mut self, bytes: u64) {
797 debug_assert!(bytes <= self.reserved);
798 self.counter
799 .fetch_sub(self.reserved.saturating_sub(bytes), Ordering::AcqRel);
800 self.committed = true;
801 }
802}
803
804impl Drop for RemoteDownloadReservation {
805 fn drop(&mut self) {
806 if !self.committed {
807 self.counter.fetch_sub(self.reserved, Ordering::AcqRel);
808 }
809 }
810}
811
812#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
813struct ExecutableIdentityKey {
814 executable: PathBuf,
815 environment: BTreeMap<String, Option<String>>,
816}
817
818const PERSISTED_EXECUTABLE_IDENTITY_VERSION: u8 = 2;
820
821#[derive(Debug, Serialize, Deserialize)]
827struct PersistedExecutableIdentity {
828 version: u8,
829 executable: PathBuf,
830 environment: BTreeMap<String, Option<String>>,
831 pins: Vec<PinnedFile>,
832 stdout: String,
833}
834
835impl CacheAgent {
836 pub fn new(cache_dir: impl Into<PathBuf>, version: impl Into<Arc<str>>) -> Self {
838 Self::build(cache_dir.into(), version.into(), None, 0)
839 }
840
841 pub fn new_remote(
843 cache_dir: impl Into<PathBuf>,
844 version: impl Into<Arc<str>>,
845 remote: AgentRemoteCache,
846 ) -> Self {
847 Self::build(
848 cache_dir.into(),
849 version.into(),
850 Some(remote),
851 DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES,
852 )
853 }
854
855 pub fn new_remote_with_download_limit(
857 cache_dir: impl Into<PathBuf>,
858 version: impl Into<Arc<str>>,
859 remote: AgentRemoteCache,
860 max_remote_download_bytes: u64,
861 ) -> Self {
862 Self::build(
863 cache_dir.into(),
864 version.into(),
865 Some(remote),
866 max_remote_download_bytes,
867 )
868 }
869
870 fn build(
871 cache_dir: PathBuf,
872 version: Arc<str>,
873 remote: Option<AgentRemoteCache>,
874 remote_download_limit: u64,
875 ) -> Self {
876 let remote_mode = remote
877 .as_ref()
878 .map_or(RemoteCacheMode::ReadOnly, |remote| remote.mode);
879 let remote_staging_dir = remote.as_ref().map_or_else(
880 || cache_dir.join("remote"),
881 |remote| remote.staging_dir.clone(),
882 );
883 let remote = remote.map(|remote| Arc::new(remote.client));
884 let stats = Arc::new(AtomicAgentStats::default());
885 let remote_transfers = Arc::new(tokio::sync::Semaphore::new(MAX_REMOTE_TRANSFERS));
886 let uploads = remote
889 .clone()
890 .filter(|_| remote_mode.writes())
891 .map(|client| {
892 UploadQueue::new(
893 client,
894 Arc::new(AgentUploadSink {
895 stats: stats.clone(),
896 }),
897 remote_transfers.clone(),
898 )
899 });
900 Self {
901 cas: LocalCas::new(cache_dir.clone()),
902 actions: LocalActionCache::new(cache_dir.clone()),
903 verified_blobs: Arc::new(Mutex::new(BTreeMap::new())),
904 version,
905 write_locks: Arc::new(Mutex::new(BTreeMap::new())),
906 action_locks: Arc::new(Mutex::new(BTreeMap::new())),
907 stats,
908 observer: None,
909 observer_emission: Arc::new(Mutex::new(())),
910 task_loader: None,
911 executable_identities: Arc::new(Mutex::new(BTreeMap::new())),
912 identity_dir: Arc::new(cache_dir.join("executable-identities").join("v1")),
913 manifest_dir: Arc::new(task_manifest_dir(&cache_dir)),
914 task_actions: Arc::new(Mutex::new(BTreeMap::new())),
915 task_fallbacks: Arc::new(Mutex::new(BTreeMap::new())),
916 next_task_run: Arc::new(AtomicU64::new(0)),
917 manifest_write_lock: Arc::new(Mutex::new(())),
918 remote,
919 remote_mode,
920 remote_staging_dir: Arc::new(remote_staging_dir),
921 remote_download_limit,
922 remote_download_bytes: Arc::new(AtomicU64::new(0)),
923 pending_remote_actions: Arc::new(Mutex::new(BTreeMap::new())),
924 remote_transfers,
925 prefetch_transfers: Arc::new(tokio::sync::Semaphore::new(MAX_PREFETCH_TRANSFERS)),
926 prefetch_tasks: Arc::new(Mutex::new(Vec::new())),
927 diagnostics: Arc::new(Mutex::new(BTreeSet::new())),
928 file_digests: Arc::new(Mutex::new(BTreeMap::new())),
929 file_digest_locks: Arc::new(Mutex::new(BTreeMap::new())),
930 file_digest_permits: Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_FILE_DIGESTS)),
931 blob_verification_permits: Arc::new(tokio::sync::Semaphore::new(
932 MAX_CONCURRENT_BLOB_VERIFICATIONS,
933 )),
934 #[cfg(test)]
935 file_digest_reads: Arc::new(AtomicU64::new(0)),
936 uploads,
937 }
938 }
939
940 #[must_use]
942 pub fn with_observer(mut self, observer: Arc<dyn AgentEventObserver>) -> Self {
943 self.observer = Some(observer);
944 self
945 }
946
947 pub fn with_task_loader(mut self, loader: Arc<TaskLoader>) -> Self {
954 self.task_loader = Some(loader);
955 self
956 }
957
958 async fn ensure_task_loaded(&self, task: &str) {
961 let Some(loader) = &self.task_loader else {
962 return;
963 };
964 if self.task_actions.lock().unwrap().contains_key(task) {
965 return;
966 }
967 loader(self, task).await;
968 }
969
970 fn emit(&self, event: impl FnOnce() -> AgentEvent) {
971 if let Some(observer) = &self.observer {
972 let _emission = self.observer_emission.lock().unwrap();
973 observer.event(event());
974 }
975 }
976
977 fn emit_action(&self, diagnostic: Option<AgentEvent>, action: AgentEvent) {
978 if let Some(observer) = &self.observer {
979 let _emission = self.observer_emission.lock().unwrap();
980 if let Some(diagnostic) = diagnostic {
981 observer.event(diagnostic);
982 }
983 observer.event(action);
984 }
985 }
986
987 fn reserve_remote_download(&self, bytes: u64) -> Result<RemoteDownloadReservation> {
988 let mut current = self.remote_download_bytes.load(Ordering::Acquire);
989 loop {
990 let next = current
991 .checked_add(bytes)
992 .ok_or_else(|| eyre::eyre!("remote cache download budget overflowed"))?;
993 if next > self.remote_download_limit {
994 bail!(
995 "remote cache download budget exceeded: {} bytes requested with {} of {} bytes already used",
996 bytes,
997 current,
998 self.remote_download_limit
999 );
1000 }
1001 match self.remote_download_bytes.compare_exchange_weak(
1002 current,
1003 next,
1004 Ordering::AcqRel,
1005 Ordering::Acquire,
1006 ) {
1007 Ok(_) => {
1008 return Ok(RemoteDownloadReservation {
1009 counter: self.remote_download_bytes.clone(),
1010 reserved: bytes,
1011 committed: false,
1012 });
1013 }
1014 Err(observed) => current = observed,
1015 }
1016 }
1017 }
1018
1019 fn reserve_remote_download_up_to(&self, requested: u64) -> Result<RemoteDownloadReservation> {
1020 let mut current = self.remote_download_bytes.load(Ordering::Acquire);
1021 loop {
1022 let reserved = requested.min(self.remote_download_limit.saturating_sub(current));
1023 let next = current + reserved;
1024 match self.remote_download_bytes.compare_exchange_weak(
1025 current,
1026 next,
1027 Ordering::AcqRel,
1028 Ordering::Acquire,
1029 ) {
1030 Ok(_) => {
1031 return Ok(RemoteDownloadReservation {
1032 counter: self.remote_download_bytes.clone(),
1033 reserved,
1034 committed: false,
1035 });
1036 }
1037 Err(observed) => current = observed,
1038 }
1039 }
1040 }
1041
1042 pub async fn begin_task(&self, task: &str) -> Result<String> {
1044 self.begin_task_with_remote_errors(task, false, None, true)
1045 .await
1046 }
1047
1048 pub async fn begin_task_on_prediction(&self, task: &str) -> Result<String> {
1051 self.begin_task_with_remote_errors(task, false, None, false)
1052 .await
1053 }
1054
1055 pub async fn begin_session_task(&self, task: &str) -> Result<()> {
1061 self.begin_task_with_remote_errors(task, false, Some(task.to_string()), false)
1062 .await?;
1063 Ok(())
1064 }
1065
1066 pub fn register_task_fallbacks(
1077 &self,
1078 task: &str,
1079 fallbacks: impl Fn() -> Vec<String> + Send + Sync + 'static,
1080 ) -> Result<()> {
1081 validate_task_identity(task)?;
1082 self.task_fallbacks
1083 .lock()
1084 .unwrap()
1085 .insert(task.to_string(), Arc::new(fallbacks));
1086 Ok(())
1087 }
1088
1089 pub async fn prefetch_task(&self, task: &str) -> Result<String> {
1091 let run = self
1092 .begin_task_with_remote_errors(task, true, None, true)
1093 .await?;
1094 self.wait_for_prefetches().await;
1095 Ok(run)
1096 }
1097
1098 async fn begin_task_with_remote_errors(
1099 &self,
1100 task: &str,
1101 strict: bool,
1102 run: Option<String>,
1103 eager_prefetch: bool,
1104 ) -> Result<String> {
1105 validate_task_identity(task)?;
1106 let early_manifest = {
1112 let _write_guard = self.manifest_write_lock.lock().unwrap();
1113 let _file_guard = self.lock_task_manifest(task)?;
1114 self.load_task_manifest(task)?
1115 };
1116 let early_actions: BTreeSet<_> = early_manifest
1117 .iter()
1118 .flat_map(|manifest| manifest.predictions.iter())
1119 .map(|prediction| prediction.action.clone())
1120 .collect();
1121 if eager_prefetch {
1122 self.spawn_prefetch_predictions(
1123 early_manifest
1124 .as_ref()
1125 .map(|manifest| manifest.predictions.clone())
1126 .unwrap_or_default(),
1127 );
1128 }
1129 let (remote_manifest, mut remote_etag, mut remote_predictions, remote_missing) =
1130 if self.remote.is_some() && self.remote_mode.reads() {
1131 match self.get_remote_task_manifest(task).await {
1132 Ok(Some((manifest, etag))) => {
1133 let predictions = manifest
1134 .predictions
1135 .iter()
1136 .map(|prediction| (prediction.invocation.clone(), prediction.clone()))
1137 .collect();
1138 (Some(manifest), Some(etag), Some(predictions), false)
1139 }
1140 Ok(None) => (None, None, Some(BTreeMap::new()), true),
1141 Err(error) => {
1142 if strict {
1143 return Err(error).wrap_err_with(|| {
1144 format!("remote task action manifest lookup failed for {task}")
1145 });
1146 }
1147 self.note_remote_failure();
1148 warn!("remote task action manifest lookup failed for {task}: {error}");
1149 (None, None, None, false)
1150 }
1151 }
1152 } else {
1153 (None, None, None, false)
1154 };
1155 let remote_publication_base = remote_manifest.clone();
1156 let manifest = if self.remote.is_some() && self.remote_mode.reads() {
1160 let _write_guard = self.manifest_write_lock.lock().unwrap();
1161 let _file_guard = self.lock_task_manifest(task)?;
1162 let local_manifest = self.load_task_manifest(task)?;
1163 let local_baseline = local_manifest.clone();
1164 let manifest = match (remote_manifest, local_manifest) {
1165 (Some(remote), Some(local)) => {
1166 let required = local
1167 .predictions
1168 .iter()
1169 .map(|prediction| prediction.invocation.clone())
1170 .collect();
1171 match merge_task_manifests(
1172 task,
1173 remote,
1174 local.clone(),
1175 &required,
1176 &BTreeSet::new(),
1177 ) {
1178 Ok(manifest) => Some(manifest),
1179 Err(error) => {
1180 warn!("remote task action manifest merge failed for {task}: {error}");
1181 remote_etag = None;
1182 Some(local)
1183 }
1184 }
1185 }
1186 (Some(remote), None) => Some(remote),
1187 (None, local) => local,
1188 };
1189 if manifest != local_baseline
1190 && let Some(manifest) = &manifest
1191 {
1192 self.persist_task_manifest(manifest)?;
1193 }
1194 manifest
1195 } else {
1196 early_manifest
1197 };
1198 let (manifest, remote_etag, inherited) = match manifest {
1199 Some(manifest) => (Some(manifest), remote_etag, false),
1200 None => {
1203 let (manifest, inherited_provenance, inherited) =
1204 self.inherit_task_manifest(task, strict).await?;
1205 if inherited {
1206 remote_predictions = inherited_provenance;
1207 }
1208 (manifest, None, inherited)
1209 }
1210 };
1211 let mut state = if let Some(manifest) = manifest {
1212 let write_only_baseline = self.remote_mode.writes()
1213 && !self.remote_mode.reads()
1214 && !manifest.predictions.is_empty();
1215 TaskActionState {
1216 manifest: task.to_string(),
1217 baseline_loaded: true,
1218 predictions: manifest
1219 .predictions
1220 .into_iter()
1221 .map(|prediction| (prediction.invocation.clone(), prediction))
1222 .collect(),
1223 pending_predictions: BTreeMap::new(),
1224 changed_predictions: BTreeMap::new(),
1225 remote_predictions,
1226 remote_manifest: remote_publication_base,
1227 remote_publication_needed: self.remote.is_some()
1228 && self.remote_mode.writes()
1229 && (remote_missing || inherited || write_only_baseline),
1230 prefetched_adapters: BTreeSet::new(),
1231 remote_etag,
1232 }
1233 } else {
1234 TaskActionState {
1235 manifest: task.to_string(),
1236 baseline_loaded: true,
1237 remote_predictions,
1238 remote_manifest: remote_publication_base,
1239 remote_etag,
1240 ..TaskActionState::default()
1241 }
1242 };
1243 let run = run.unwrap_or_else(|| {
1244 let sequence = self.next_task_run.fetch_add(1, Ordering::Relaxed);
1245 CacheDigest::blake3(format!("{task}\0{}\0{sequence}", std::process::id()).as_bytes())
1246 .hash
1247 });
1248 self.stats.predictions_loaded.fetch_max(
1252 state.predictions.len().try_into().unwrap_or(u64::MAX),
1253 Ordering::Relaxed,
1254 );
1255 if eager_prefetch {
1256 state.prefetched_adapters.extend(
1257 state
1258 .predictions
1259 .values()
1260 .map(|prediction| prediction.adapter.clone()),
1261 );
1262 }
1263 let predictions = state
1267 .predictions
1268 .values()
1269 .filter(|prediction| !early_actions.contains(&prediction.action))
1270 .cloned()
1271 .collect();
1272 self.task_actions.lock().unwrap().insert(run.clone(), state);
1273 if eager_prefetch {
1274 self.spawn_prefetch_predictions(predictions);
1275 }
1276 Ok(run)
1277 }
1278
1279 async fn inherit_task_manifest(
1287 &self,
1288 task: &str,
1289 strict: bool,
1290 ) -> Result<(
1291 Option<TaskActionManifest>,
1292 Option<BTreeMap<CacheDigest, ActionPrediction>>,
1293 bool,
1294 )> {
1295 let fallbacks = self.task_fallbacks.lock().unwrap().get(task).cloned();
1296 let Some(fallbacks) = fallbacks else {
1297 return Ok((None, None, false));
1298 };
1299 let named = tokio::task::spawn_blocking(move || fallbacks()).await?;
1300 let mut tried = BTreeSet::from([task.to_string()]);
1301 let candidates = named.into_iter().map(|identity| (identity, true)).chain(
1302 self.newest_task_identities()
1303 .into_iter()
1304 .map(|identity| (identity, false)),
1305 );
1306 for (identity, named) in candidates {
1307 if !tried.insert(identity.clone()) || validate_task_identity(&identity).is_err() {
1308 continue;
1309 }
1310 let mut found = self.load_task_manifest(&identity)?;
1311 let mut provenance = None;
1312 if self.remote_mode.reads() && (found.is_some() || named) {
1313 match self.get_remote_task_manifest(&identity).await {
1314 Ok(Some((remote, _))) => {
1315 provenance = Some(
1316 remote
1317 .predictions
1318 .iter()
1319 .map(|prediction| {
1320 (prediction.invocation.clone(), prediction.clone())
1321 })
1322 .collect(),
1323 );
1324 if found.is_none() {
1325 found = Some(remote);
1326 }
1327 }
1328 Ok(None) => provenance = Some(BTreeMap::new()),
1329 Err(error) => {
1330 if strict && found.is_none() {
1335 return Err(error).wrap_err_with(|| {
1336 format!("remote task action manifest lookup failed for {identity}")
1337 });
1338 }
1339 self.note_remote_failure();
1340 warn!("remote task action manifest lookup failed for {identity}: {error}");
1341 }
1342 }
1343 }
1344 let Some(mut manifest) = found else {
1345 continue;
1346 };
1347 if manifest.predictions.is_empty()
1348 || (!named && manifest.predictions.len() > MAX_TASK_ACTION_PREDICTIONS / 2)
1349 {
1350 continue;
1351 }
1352 debug!(
1353 "no predictions were recorded for {task}; inheriting {} from {identity}",
1354 manifest.predictions.len()
1355 );
1356 manifest.task = task.to_string();
1357 validate_task_manifest(&manifest, task)?;
1358 let _write_guard = self.manifest_write_lock.lock().unwrap();
1359 let _file_guard = self.lock_task_manifest(task)?;
1360 if let Some(recorded) = self.load_task_manifest(task)? {
1364 return Ok((Some(recorded), None, false));
1365 }
1366 self.persist_task_manifest(&manifest)?;
1367 return Ok((Some(manifest), provenance, true));
1368 }
1369 Ok((None, None, false))
1370 }
1371
1372 fn newest_task_identities(&self) -> Vec<String> {
1374 let Ok(entries) = fs::read_dir(self.manifest_dir.as_path()) else {
1375 return Vec::new();
1376 };
1377 let mut manifests: Vec<(std::time::SystemTime, String)> = entries
1378 .filter_map(|entry| {
1379 let entry = entry.ok()?;
1380 let name = entry.file_name();
1381 let identity = name.to_str()?.strip_suffix(".json")?.to_string();
1382 validate_task_identity(&identity).ok()?;
1383 let modified = entry.metadata().ok()?.modified().ok()?;
1384 Some((modified, identity))
1385 })
1386 .collect();
1387 manifests.sort_by(|left, right| right.cmp(left));
1388 manifests
1389 .into_iter()
1390 .take(NEWEST_MANIFEST_CANDIDATES)
1391 .map(|(_, identity)| identity)
1392 .collect()
1393 }
1394
1395 pub async fn cancel_prefetches(&self) {
1397 let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
1398 for task in &tasks {
1399 task.abort();
1400 }
1401 for task in tasks {
1402 if let Err(error) = task.await
1403 && !error.is_cancelled()
1404 {
1405 warn!("remote action prefetch task failed: {error}");
1406 }
1407 }
1408 }
1409
1410 pub async fn wait_for_uploads(&self) {
1417 let Some(uploads) = &self.uploads else {
1418 return;
1419 };
1420 let _timer = AtomicDurationTimer::start(&self.stats.upload_drain_duration_ns);
1421 uploads.drain().await;
1422 }
1423
1424 async fn wait_for_prefetches(&self) {
1425 let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
1426 for task in tasks {
1427 if let Err(error) = task.await {
1428 warn!("remote action prefetch task failed: {error}");
1429 }
1430 }
1431 }
1432
1433 pub async fn commit_task(&self, run: &str) -> Result<()> {
1435 self.commit_task_actions(run).await.map(|_| ())
1436 }
1437
1438 pub async fn commit_task_actions(&self, run: &str) -> Result<Vec<ActionPrediction>> {
1444 validate_task_identity(run)?;
1445 let mut state = {
1446 let mut runs = self.task_actions.lock().unwrap();
1447 let state = runs
1448 .get(run)
1449 .ok_or_else(|| eyre::eyre!("task action manifest baseline was not loaded"))?;
1450 if !state.baseline_loaded {
1451 bail!("task action manifest baseline was not loaded");
1452 }
1453 debug!(
1458 "committing {} changed of {} recorded predictions for {}",
1459 state.changed_predictions.len(),
1460 state.pending_predictions.len(),
1461 state.manifest
1462 );
1463 let refresh_hit_order = should_refresh_task_prediction_order(
1464 state.predictions.len(),
1465 !state.pending_predictions.is_empty(),
1466 );
1467 if state.changed_predictions.is_empty()
1468 && !refresh_hit_order
1469 && !state.remote_publication_needed
1470 {
1471 let completed = state.pending_predictions.values().cloned().collect();
1472 runs.remove(run);
1473 return Ok(completed);
1474 }
1475 state.clone()
1476 };
1477 let task = state.manifest.clone();
1478 validate_task_identity(&task)?;
1479 let write_only = self.remote_mode.writes() && !self.remote_mode.reads();
1480 let write_only_remote = if write_only {
1486 match self.get_remote_task_manifest(&task).await {
1487 Ok(Some((manifest, etag))) => {
1488 let remote_predictions: BTreeMap<_, _> = manifest
1489 .predictions
1490 .iter()
1491 .map(|prediction| (prediction.invocation.clone(), prediction.clone()))
1492 .collect();
1493 let has_missing_fallback = state
1494 .predictions
1495 .keys()
1496 .any(|invocation| !remote_predictions.contains_key(invocation));
1497 state.remote_predictions = if has_missing_fallback {
1498 None
1499 } else {
1500 Some(remote_predictions)
1501 };
1502 state.remote_etag = Some(etag);
1503 state.remote_manifest = Some(manifest.clone());
1504 state.remote_publication_needed = has_missing_fallback;
1509 Some(manifest)
1510 }
1511 Ok(None) => {
1512 state.remote_predictions = if state.predictions.is_empty() {
1513 Some(BTreeMap::new())
1514 } else {
1515 None
1516 };
1517 state.remote_publication_needed = !state.predictions.is_empty();
1518 None
1519 }
1520 Err(error) => {
1521 state.remote_predictions = None;
1522 state.remote_publication_needed = false;
1523 self.note_remote_failure();
1524 warn!("remote task action manifest lookup failed for {task}: {error}");
1525 None
1526 }
1527 }
1528 } else {
1529 None
1530 };
1531 let completed = state
1532 .pending_predictions
1533 .values()
1534 .cloned()
1535 .collect::<Vec<_>>();
1536 if write_only
1537 && state.changed_predictions.is_empty()
1538 && !should_refresh_task_prediction_order(
1539 state.predictions.len(),
1540 !state.pending_predictions.is_empty(),
1541 )
1542 && !state.remote_publication_needed
1543 {
1544 self.task_actions.lock().unwrap().remove(run);
1545 return Ok(completed);
1546 }
1547 let (manifest, concurrent_updates) = {
1548 let _write_guard = self.manifest_write_lock.lock().unwrap();
1549 let _file_guard = self.lock_task_manifest(&task)?;
1550 let current_manifest = self.load_task_manifest(&task)?;
1551 let predictions = current_manifest
1552 .as_ref()
1553 .map(|manifest| manifest.predictions.clone())
1554 .unwrap_or_default();
1555 let concurrent_updates = concurrent_task_predictions(&predictions, &state);
1556 let effective_updates = effective_task_updates(&predictions, &state);
1557 let required =
1558 required_local_predictions(&predictions, &effective_updates, &concurrent_updates);
1559 let predictions = update_task_predictions(predictions, &effective_updates, &required)?;
1560 let manifest = TaskActionManifest {
1561 version: TASK_ACTION_MANIFEST_VERSION,
1562 task: task.clone(),
1563 predictions,
1564 };
1565 validate_task_manifest(&manifest, &task)?;
1566 if current_manifest.as_ref() != Some(&manifest) {
1567 self.persist_task_manifest(&manifest)?;
1568 }
1569 (manifest, concurrent_updates)
1570 };
1571 self.task_actions.lock().unwrap().remove(run);
1572 if self.remote_mode.writes() {
1573 let mut manifest = manifest;
1577 let (held, failed) = if let Some(uploads) = &self.uploads {
1578 let actions: Vec<CacheDigest> = manifest
1579 .predictions
1580 .iter()
1581 .map(|prediction| prediction.action.clone())
1582 .collect();
1583 uploads.wait_for_action_outcomes(&actions).await
1584 } else {
1585 (BTreeSet::new(), BTreeSet::new())
1586 };
1587 let proven_uploads = held.difference(&failed).cloned().collect();
1588 let (withhold_candidates, replaced, unpublished) =
1589 if let Some(remote_predictions) = &state.remote_predictions {
1590 let unproven_predictions =
1591 unproven_task_predictions(&manifest.predictions, remote_predictions);
1592 let unproven = unproven_predictions.values().cloned().collect();
1593 let candidates = unproven_predictions.keys().cloned().collect();
1594 let replaced = replaced_task_predictions(remote_predictions, &candidates);
1595 let unpublished = unpublished_task_actions(unproven, &held, failed);
1598 (candidates, replaced, unpublished)
1599 } else {
1600 let candidates: BTreeSet<_> = state
1605 .changed_predictions
1606 .keys()
1607 .filter(|invocation| {
1608 manifest
1609 .predictions
1610 .iter()
1611 .any(|prediction| &prediction.invocation == *invocation)
1612 })
1613 .cloned()
1614 .collect();
1615 let unproven = manifest
1616 .predictions
1617 .iter()
1618 .filter(|prediction| candidates.contains(&prediction.invocation))
1619 .map(|prediction| prediction.action.clone())
1620 .collect();
1621 let unpublished = unpublished_task_actions(unproven, &held, failed);
1622 (candidates, Vec::new(), unpublished)
1623 };
1624 let prediction_count = manifest.predictions.len();
1625 let withheld = withhold_unpublished_updates(
1626 &mut manifest,
1627 &withhold_candidates,
1628 &replaced,
1629 &unpublished,
1630 );
1631 if !withheld.is_empty() {
1632 warn!(
1635 "{} of {} predicted actions were not published, so the remote task action manifest omits them",
1636 withheld.len(),
1637 prediction_count
1638 );
1639 }
1640 let published: BTreeMap<_, _> = manifest
1641 .predictions
1642 .iter()
1643 .map(|prediction| (prediction.invocation.clone(), prediction))
1644 .collect();
1645 let (required, fallbacks) = remote_publication_precedence(
1646 &state,
1647 &completed,
1648 &concurrent_updates,
1649 &proven_uploads,
1650 &published,
1651 &withheld,
1652 );
1653 let publication_base = write_only_remote.or(state.remote_manifest);
1654 if let Some(remote) = publication_base {
1655 manifest = merge_task_manifests(&task, remote, manifest, &required, &fallbacks)?;
1656 }
1657 match self
1658 .put_remote_task_manifest(&task, manifest, state.remote_etag, &required, &fallbacks)
1659 .await
1660 {
1661 Ok(remote_manifest) => {
1662 let _write_guard = self.manifest_write_lock.lock().unwrap();
1663 let reconciliation = (|| {
1664 let _file_guard = self.lock_task_manifest(&task)?;
1665 let manifest = match self.load_task_manifest(&task)? {
1666 Some(local) => {
1667 let required = local
1668 .predictions
1669 .iter()
1670 .map(|prediction| prediction.invocation.clone())
1671 .collect();
1672 match merge_task_manifests(
1673 &task,
1674 remote_manifest,
1675 local.clone(),
1676 &required,
1677 &BTreeSet::new(),
1678 ) {
1679 Ok(manifest) => manifest,
1680 Err(error) => {
1681 warn!(
1682 "remote task action manifest merge failed for {task}: {error}"
1683 );
1684 local
1685 }
1686 }
1687 }
1688 None => remote_manifest,
1689 };
1690 self.persist_task_manifest(&manifest)
1691 })();
1692 if let Err(error) = reconciliation {
1693 warn!(
1694 "remote task action manifest reconciliation failed for {task}: {error}"
1695 );
1696 }
1697 }
1698 Err(error) => {
1699 self.note_remote_failure();
1700 warn!("remote task action manifest upload failed for {task}: {error}");
1701 }
1702 }
1703 }
1704 Ok(completed)
1705 }
1706
1707 fn task_manifest_path(&self, task: &str) -> PathBuf {
1708 self.manifest_dir.join(format!("{task}.json"))
1709 }
1710
1711 fn task_manifest_lock_path(&self, task: &str) -> PathBuf {
1712 self.manifest_dir.join("locks").join(format!("{task}.lock"))
1713 }
1714
1715 fn lock_task_manifest(&self, task: &str) -> Result<fslock::LockFile> {
1716 let path = self.task_manifest_lock_path(task);
1717 fs::create_dir_all(path.parent().expect("task manifest lock has a parent"))?;
1718 let mut lock = fslock::LockFile::open(&path)?;
1719 lock.lock()?;
1720 Ok(lock)
1721 }
1722
1723 fn load_task_manifest(&self, task: &str) -> Result<Option<TaskActionManifest>> {
1724 match fs::read(self.task_manifest_path(task)) {
1725 Ok(contents) => Ok(Some(self.parse_task_manifest(task, &contents, false)?)),
1726 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
1727 Err(error) => Err(error.into()),
1728 }
1729 }
1730
1731 fn parse_task_manifest(
1732 &self,
1733 task: &str,
1734 contents: &[u8],
1735 require_canonical: bool,
1736 ) -> Result<TaskActionManifest> {
1737 let manifest: TaskActionManifest = serde_json::from_slice(contents)?;
1738 validate_task_manifest(&manifest, task)?;
1739 if require_canonical && canonical_json(&manifest)? != contents {
1740 bail!("task action manifest is not canonical JSON");
1741 }
1742 Ok(manifest)
1743 }
1744
1745 fn task_manifest_selector(task: &str) -> Result<(Vec<u8>, CacheDigest)> {
1746 TaskActionManifest::selector(task)
1747 }
1748
1749 fn persist_task_manifest(&self, manifest: &TaskActionManifest) -> Result<()> {
1750 let bytes = canonical_json(manifest)?;
1751 fs::create_dir_all(self.manifest_dir.as_path())?;
1752 let mut temporary = tempfile::NamedTempFile::new_in(self.manifest_dir.as_path())?;
1753 std::io::Write::write_all(temporary.as_file_mut(), &bytes)?;
1754 temporary.as_file_mut().sync_all()?;
1755 temporary
1756 .persist(self.task_manifest_path(&manifest.task))
1757 .map_err(|error| error.error)?;
1758 Ok(())
1759 }
1760
1761 async fn get_remote_task_manifest(
1762 &self,
1763 task: &str,
1764 ) -> Result<Option<(TaskActionManifest, String)>> {
1765 let Some(remote) = &self.remote else {
1766 return Ok(None);
1767 };
1768 let (_, selector) = Self::task_manifest_selector(task)?;
1769 let _permit = self.remote_transfers.acquire().await?;
1770 self.stats
1771 .remote_manifest_lookups
1772 .fetch_add(1, Ordering::Relaxed);
1773 let _timer = AtomicDurationTimer::start(&self.stats.remote_manifest_lookup_duration_ns);
1774 let Some(remote_manifest) = remote.get_action_manifest(&selector).await? else {
1775 return Ok(None);
1776 };
1777 let manifest = self.parse_task_manifest(task, &remote_manifest.bytes, true)?;
1778 Ok(Some((manifest, remote_manifest.etag)))
1779 }
1780
1781 async fn put_remote_task_manifest(
1782 &self,
1783 task: &str,
1784 mut manifest: TaskActionManifest,
1785 mut expected_etag: Option<String>,
1786 required: &BTreeSet<CacheDigest>,
1787 fallbacks: &BTreeSet<CacheDigest>,
1788 ) -> Result<TaskActionManifest> {
1789 let Some(remote) = &self.remote else {
1790 return Ok(manifest);
1791 };
1792 let (_, selector) = Self::task_manifest_selector(task)?;
1793 for _ in 0..4 {
1794 let bytes = canonical_json(&manifest)?;
1795 let outcome = {
1796 let _permit = self.remote_transfers.acquire().await?;
1797 remote
1798 .put_action_manifest(&selector, &bytes, expected_etag.as_deref())
1799 .await?
1800 };
1801 match outcome {
1802 ManifestPutOutcome::Stored => return Ok(manifest),
1803 ManifestPutOutcome::PreconditionFailed => {
1804 let Some((remote_manifest, etag)) = self.get_remote_task_manifest(task).await?
1805 else {
1806 expected_etag = None;
1807 continue;
1808 };
1809 manifest =
1810 merge_task_manifests(task, remote_manifest, manifest, required, fallbacks)?;
1811 expected_etag = Some(etag);
1812 }
1813 }
1814 }
1815 bail!("remote task action manifest changed too frequently")
1816 }
1817
1818 pub fn stats(&self) -> AgentStats {
1820 AgentStats {
1821 session_duration_ns: 0,
1822 lookups: self.stats.lookups.load(Ordering::Relaxed),
1823 unconsulted: self.stats.unconsulted.load(Ordering::Relaxed),
1824 incremental_compilations: self.stats.incremental_compilations.load(Ordering::Relaxed),
1825 hits: self.stats.hits.load(Ordering::Relaxed),
1826 stores: self.stats.stores.load(Ordering::Relaxed),
1827 stored_bytes: self.stats.stored_bytes.load(Ordering::Relaxed),
1828 verifications: self.stats.verifications.load(Ordering::Relaxed),
1829 divergences: self.stats.divergences.load(Ordering::Relaxed),
1830 downloaded_bytes: self.stats.downloaded_bytes.load(Ordering::Relaxed),
1831 uploaded_bytes: self.stats.uploaded_bytes.load(Ordering::Relaxed),
1832 background_uploads: self.stats.background_uploads.load(Ordering::Relaxed),
1833 background_upload_failures: self
1834 .stats
1835 .background_upload_failures
1836 .load(Ordering::Relaxed),
1837 remote_blob_pack_uploads: self.stats.remote_blob_pack_uploads.load(Ordering::Relaxed),
1838 remote_blob_pack_upload_blobs: self
1839 .stats
1840 .remote_blob_pack_upload_blobs
1841 .load(Ordering::Relaxed),
1842 upload_drain_duration_ns: self.stats.upload_drain_duration_ns.load(Ordering::Relaxed),
1843 prefetched_actions: self.stats.prefetched_actions.load(Ordering::Relaxed),
1844 predictions_loaded: self.stats.predictions_loaded.load(Ordering::Relaxed),
1845 bypasses: self.stats.bypasses.lock().unwrap().clone(),
1846 avoided_compiler_duration_ns: self
1847 .stats
1848 .avoided_compiler_duration_ns
1849 .load(Ordering::Relaxed),
1850 wrapper_phases_ns: self.stats.wrapper_phases_ns.lock().unwrap().clone(),
1851 compiler: self.stats.compiler.lock().unwrap().clone(),
1852 slow_compilations: self.stats.slow_compilations.lock().unwrap().clone(),
1853 unit_outcomes: self.stats.unit_outcomes.lock().unwrap().clone(),
1854 remote_failures: self.stats.remote_failures.load(Ordering::Relaxed),
1855 remote_manifest_lookups: self.stats.remote_manifest_lookups.load(Ordering::Relaxed),
1856 remote_manifest_lookup_duration_ns: self
1857 .stats
1858 .remote_manifest_lookup_duration_ns
1859 .load(Ordering::Relaxed),
1860 remote_action_lookups: self.stats.remote_action_lookups.load(Ordering::Relaxed),
1861 remote_action_lookup_duration_ns: self
1862 .stats
1863 .remote_action_lookup_duration_ns
1864 .load(Ordering::Relaxed),
1865 remote_blob_requests: self.stats.remote_blob_requests.load(Ordering::Relaxed),
1866 remote_blob_pack_requests: self.stats.remote_blob_pack_requests.load(Ordering::Relaxed),
1867 remote_blob_pack_blobs: self.stats.remote_blob_pack_blobs.load(Ordering::Relaxed),
1868 remote_blob_transfer_duration_ns: self
1869 .stats
1870 .remote_blob_transfer_duration_ns
1871 .load(Ordering::Relaxed),
1872 local_cas_write_duration_ns: self
1873 .stats
1874 .local_cas_write_duration_ns
1875 .load(Ordering::Relaxed),
1876 prefetch_runs: self.stats.prefetch_runs.load(Ordering::Relaxed),
1877 prefetch_duration_ns: self.stats.prefetch_duration_ns.load(Ordering::Relaxed),
1878 materialization_duration_ns: self
1879 .stats
1880 .materialization_duration_ns
1881 .load(Ordering::Relaxed),
1882 restored_output_files: self.stats.restored_output_files.load(Ordering::Relaxed),
1883 restored_output_bytes: self.stats.restored_output_bytes.load(Ordering::Relaxed),
1884 reflinked_output_files: self.stats.reflinked_output_files.load(Ordering::Relaxed),
1885 reflinked_output_bytes: self.stats.reflinked_output_bytes.load(Ordering::Relaxed),
1886 hardlinked_output_files: self.stats.hardlinked_output_files.load(Ordering::Relaxed),
1887 hardlinked_output_bytes: self.stats.hardlinked_output_bytes.load(Ordering::Relaxed),
1888 copied_output_files: self.stats.copied_output_files.load(Ordering::Relaxed),
1889 copied_output_bytes: self.stats.copied_output_bytes.load(Ordering::Relaxed),
1890 reused_output_files: self.stats.reused_output_files.load(Ordering::Relaxed),
1891 reused_output_bytes: self.stats.reused_output_bytes.load(Ordering::Relaxed),
1892 }
1893 }
1894
1895 pub async fn handle_requests(
1902 &self,
1903 requests: impl IntoIterator<Item = AgentRequest>,
1904 ) -> Vec<AgentResponse> {
1905 let mut connection = ConnectionUploads::default();
1906 let mut responses = Vec::new();
1907 for request in requests {
1908 responses.push(self.respond_on(request, &mut connection).await);
1909 }
1910 self.release_bypass(&mut connection);
1911 responses
1912 }
1913
1914 fn release_bypass(&self, connection: &mut ConnectionUploads) {
1916 if let Some(kind) = connection.take_bypass() {
1917 self.emit(|| AgentEvent::Bypass { kind });
1918 }
1919 }
1920
1921 fn write_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1922 Self::digest_lock(&self.write_locks, digest)
1923 }
1924
1925 fn action_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1926 Self::digest_lock(&self.action_locks, digest)
1927 }
1928
1929 fn digest_lock(
1930 locks: &Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>,
1931 digest: &CacheDigest,
1932 ) -> Arc<tokio::sync::Mutex<()>> {
1933 let mut locks = locks.lock().unwrap();
1934 locks.retain(|_, lock| lock.strong_count() > 0);
1935 if let Some(lock) = locks.get(digest).and_then(Weak::upgrade) {
1936 return lock;
1937 }
1938 let lock = Arc::new(tokio::sync::Mutex::new(()));
1939 locks.insert(digest.clone(), Arc::downgrade(&lock));
1940 lock
1941 }
1942
1943 #[cfg(test)]
1948 async fn respond(&self, request: AgentRequest) -> AgentResponse {
1949 let mut connection = ConnectionUploads::default();
1950 let response = self.respond_on(request, &mut connection).await;
1951 self.release_bypass(&mut connection);
1952 response
1953 }
1954
1955 async fn respond_on(
1956 &self,
1957 request: AgentRequest,
1958 connection: &mut ConnectionUploads,
1959 ) -> AgentResponse {
1960 match &request {
1961 AgentRequest::FindActionPrediction { task, .. }
1962 | AgentRequest::RecordActionPrediction { task, .. } => {
1963 self.ensure_task_loaded(task).await;
1964 }
1965 _ => {}
1966 }
1967 let result = match request {
1968 AgentRequest::BeginTask { task } => self
1969 .begin_task(&task)
1970 .await
1971 .map(|run| AgentResponse::TaskBegun { run }),
1972 AgentRequest::CommitTask { run } => self
1973 .commit_task(&run)
1974 .await
1975 .map(|()| AgentResponse::TaskCommitted),
1976 AgentRequest::FindBlob { digest } => self.find_blob(&digest).await,
1977 AgentRequest::FindBlobs { digests } => self.find_blobs(digests).await,
1978 AgentRequest::StoreBlob { digest, source } => {
1979 self.store_blob(&digest, &source, connection).await
1980 }
1981 AgentRequest::FindActionResult { action } => self.find_action_result(&action).await,
1988 AgentRequest::RecordActionHit {
1989 action,
1990 restore,
1991 crate_name,
1992 } => {
1993 let diagnostic = connection
1994 .take_action_diagnostic("hit", crate_name.as_deref())
1995 .map(|diagnostic| AgentEvent::ActionDiagnostic {
1996 outcome: "hit".into(),
1997 crate_name: crate_name.clone(),
1998 diagnostic,
1999 });
2000 self.record_action_hit(&action, restore, crate_name, diagnostic)
2001 }
2002 AgentRequest::RecordBypass { kind } => {
2003 *self
2004 .stats
2005 .bypasses
2006 .lock()
2007 .unwrap()
2008 .entry(kind.clone())
2009 .or_insert(0) += 1;
2010 if let Some(unclaimed) = connection.hold_bypass(kind) {
2014 self.emit(|| AgentEvent::Bypass { kind: unclaimed });
2015 }
2016 Ok(AgentResponse::BypassRecorded)
2017 }
2018 AgentRequest::RecordUnconsulted => {
2019 self.stats.unconsulted.fetch_add(1, Ordering::Relaxed);
2020 self.emit(|| AgentEvent::Unconsulted);
2021 Ok(AgentResponse::UnconsultedRecorded)
2022 }
2023 AgentRequest::RecordWarning { message } => match parse_action_diagnostic(&message) {
2024 Some(Ok((outcome, crate_name, diagnostic))) => {
2025 connection.record_action_diagnostic(outcome, crate_name, diagnostic);
2026 Ok(AgentResponse::WarningRecorded)
2027 }
2028 Some(Err(error)) => Err(error),
2029 None => self.record_diagnostic(DiagnosticLevel::Warning, message),
2030 },
2031 AgentRequest::RecordDebug { target, message } => (|| {
2032 if target.is_empty()
2033 || target.len() > 256
2034 || target.contains(['\n', '\r', '\0'])
2035 || message.is_empty()
2036 || message.len() > MAX_DIAGNOSTIC_BYTES
2037 || message.contains(['\n', '\r', '\0'])
2038 {
2039 bail!("invalid shim debug record");
2040 }
2041 if target == "mbx::unit-outcome" {
2042 let (unit, outcome): (String, String) = serde_json::from_str(&message)?;
2043 let (outcome, _) = lookup_outcome(&outcome);
2047 if unit.len() > 256
2048 || !unit.contains(':')
2049 || !matches!(
2050 outcome,
2051 "hit" | "miss" | "bypass" | "unconsulted" | "verification"
2052 )
2053 {
2054 bail!("invalid unit outcome");
2055 }
2056 self.stats
2057 .unit_outcomes
2058 .lock()
2059 .unwrap()
2060 .entry(unit)
2061 .or_default()
2062 .insert(outcome.to_string());
2063 } else {
2064 debug!(target: &target, "{message}");
2065 }
2066 Ok(AgentResponse::DebugRecorded)
2067 })(),
2068 AgentRequest::RecordError { message } => {
2069 self.record_diagnostic(DiagnosticLevel::Error, message)
2070 }
2071 AgentRequest::FindFileDigests { scope, files } => self.find_file_digests(scope, files),
2072 AgentRequest::JoinActionPromise {
2073 adapter,
2074 invocation,
2075 } => self.join_action_promise(&adapter, &invocation).await,
2076 AgentRequest::CompleteActionPromise { claim, prediction } => {
2077 self.complete_action_promise(&claim, &prediction).await
2078 }
2079 AgentRequest::ResolveFileDigests { scope, files } => {
2080 self.resolve_file_digests(scope, files).await
2081 }
2082 AgentRequest::RecordFileDigests { scope, entries } => {
2083 self.record_file_digests(scope, entries)
2084 }
2085 AgentRequest::RecordWrapperTiming { timing } => (|| {
2086 if timing.spans.len() > 512
2087 || timing.phases_ns.len() > 32
2088 || timing.adapter.len() > 64
2089 || timing.unit.as_ref().is_some_and(|s| s.len() > 1024)
2090 || timing.unit_id.as_ref().is_some_and(|s| s.len() > 256)
2091 || timing.dependencies.len() > 4096
2092 || timing.dependencies.iter().any(|s| s.len() > 256)
2093 || timing.phases_ns.keys().any(|s| s.len() > 64)
2094 || timing.spans.iter().any(|s| s.name.len() > 64)
2095 {
2096 bail!("wrapper timing exceeds telemetry limits");
2097 }
2098 {
2099 let mut phases = self.stats.wrapper_phases_ns.lock().unwrap();
2100 for (phase, duration) in &timing.phases_ns {
2101 let total = phases.entry(phase.clone()).or_default();
2102 *total = total.saturating_add(*duration);
2103 }
2104 }
2105 self.emit(|| AgentEvent::WrapperTiming { timing });
2106 Ok(AgentResponse::WrapperTimingRecorded)
2107 })(),
2108 AgentRequest::RecordCompilerInvocation {
2109 outcome,
2110 crate_name,
2111 duration_ns,
2112 } => {
2113 let (outcome, incremental) = lookup_outcome(&outcome);
2117 let outcome = outcome.to_string();
2118 let bypass = if outcome == "bypass" {
2119 connection.take_bypass()
2120 } else {
2121 self.release_bypass(connection);
2122 None
2123 };
2124 let diagnostic = connection
2125 .take_action_diagnostic(&outcome, crate_name.as_deref())
2126 .map(|diagnostic| AgentEvent::ActionDiagnostic {
2127 outcome: outcome.clone(),
2128 crate_name: crate_name.clone(),
2129 diagnostic,
2130 });
2131 self.record_compiler_invocation(
2132 &outcome,
2133 incremental,
2134 crate_name.as_deref(),
2135 duration_ns,
2136 diagnostic,
2137 bypass,
2138 )
2139 }
2140 AgentRequest::RecordActionVerification { matched, restore } => {
2141 self.record_materialization(restore);
2142 self.stats.verifications.fetch_add(1, Ordering::Relaxed);
2143 if !matched {
2144 self.stats.divergences.fetch_add(1, Ordering::Relaxed);
2145 }
2146 self.emit(|| AgentEvent::Verification { matched, restore });
2147 Ok(AgentResponse::ActionVerificationRecorded)
2148 }
2149 AgentRequest::StoreActionResult { result } => {
2150 self.store_action_result(&result, connection).await
2151 }
2152 AgentRequest::FindActionPrediction { task, invocation } => {
2153 self.find_action_prediction(&task, &invocation)
2154 }
2155 AgentRequest::RecordActionPrediction { task, prediction } => {
2156 self.record_action_prediction(&task, prediction)
2157 }
2158 AgentRequest::FindExecutableIdentity {
2159 executable,
2160 environment,
2161 } => self.find_executable_identity(executable, environment),
2162 AgentRequest::StoreExecutableIdentity {
2163 executable,
2164 environment,
2165 stdout,
2166 pins,
2167 } => self.store_executable_identity(executable, environment, stdout, pins),
2168 AgentRequest::Hello { .. } => {
2169 Err(eyre::eyre!("hello is only valid as the first request"))
2170 }
2171 };
2172 result.unwrap_or_else(|error| AgentResponse::Error {
2173 message: error.to_string(),
2174 })
2175 }
2176
2177 async fn find_blob(&self, digest: &CacheDigest) -> Result<AgentResponse> {
2178 if let Some(path) = self.find_verified_blob(digest)? {
2179 return Ok(AgentResponse::Blob { path: Some(path) });
2180 }
2181 if !self.remote_mode.reads() {
2182 return Ok(AgentResponse::Blob { path: None });
2183 }
2184 let Some(remote) = &self.remote else {
2185 return Ok(AgentResponse::Blob { path: None });
2186 };
2187 match self.fetch_remote_blob(remote, digest).await {
2188 Ok(path) => Ok(AgentResponse::Blob { path: Some(path) }),
2189 Err(error) => {
2190 warn!(
2191 "remote cache blob lookup failed for {}: {error}",
2192 digest.hash
2193 );
2194 Ok(AgentResponse::Blob { path: None })
2195 }
2196 }
2197 }
2198
2199 async fn find_blobs(&self, digests: Vec<CacheDigest>) -> Result<AgentResponse> {
2200 let mut paths = BTreeMap::new();
2201 let mut cold = BTreeSet::new();
2202 for digest in &digests {
2203 match self.find_remembered_blob(digest) {
2204 Some(path) => {
2205 paths.insert(digest.clone(), path);
2206 }
2207 None => {
2208 cold.insert(digest.clone());
2209 }
2210 }
2211 }
2212 let verified = stream::iter(cold)
2217 .map(|digest| self.verify_cold_blob(digest))
2218 .buffered(MAX_CONCURRENT_BLOB_VERIFICATIONS)
2219 .collect::<Vec<_>>()
2220 .await;
2221 let mut missing = Vec::new();
2222 for verified in verified {
2223 match verified? {
2224 (digest, Some(path)) => {
2225 paths.insert(digest, path);
2226 }
2227 (digest, None) => missing.push(digest),
2228 }
2229 }
2230
2231 if !missing.is_empty()
2232 && self.remote_mode.reads()
2233 && let Some(remote) = &self.remote
2234 {
2235 paths.extend(self.fetch_remote_blobs(remote, missing, None).await);
2236 }
2237
2238 Ok(AgentResponse::Blobs {
2239 paths: digests
2240 .into_iter()
2241 .map(|digest| paths.get(&digest).cloned())
2242 .collect(),
2243 })
2244 }
2245
2246 async fn store_blob(
2247 &self,
2248 digest: &CacheDigest,
2249 source: &Path,
2250 connection: &mut ConnectionUploads,
2251 ) -> Result<AgentResponse> {
2252 let path = {
2253 let lock = self.write_lock(digest);
2254 let _guard = lock.lock().await;
2255 if let Some(path) = self.find_verified_blob(digest)? {
2256 path
2257 } else {
2258 let timer = AtomicDurationTimer::start(&self.stats.local_cas_write_duration_ns);
2259 let path = self.cas.store_file(digest, source)?;
2260 drop(timer);
2261 self.remember_verified_blob(digest, &path);
2262 self.stats.stores.fetch_add(1, Ordering::Relaxed);
2263 self.stats
2264 .stored_bytes
2265 .fetch_add(digest.size, Ordering::Relaxed);
2266 path
2267 }
2268 };
2269 if let Some(uploads) = &self.uploads {
2272 uploads.queue_blob(digest, path.clone(), connection);
2273 }
2274 Ok(AgentResponse::Stored { path })
2275 }
2276
2277 fn find_verified_blob(&self, digest: &CacheDigest) -> Result<Option<PathBuf>> {
2278 if let Some(path) = self.find_remembered_blob(digest) {
2279 return Ok(Some(path));
2280 }
2281 let path = find_stored_blob(&self.cas, digest);
2282 if let Some(path) = &path {
2283 self.remember_verified_blob(digest, path);
2284 }
2285 Ok(path)
2286 }
2287
2288 async fn verify_cold_blob(
2292 &self,
2293 digest: CacheDigest,
2294 ) -> Result<(CacheDigest, Option<PathBuf>)> {
2295 let _permit = self.blob_verification_permits.acquire().await?;
2296 let cas = self.cas.clone();
2297 let lookup = digest.clone();
2298 let path = tokio::task::spawn_blocking(move || find_stored_blob(&cas, &lookup)).await?;
2299 if let Some(path) = &path {
2300 self.remember_verified_blob(&digest, path);
2301 }
2302 Ok((digest, path))
2303 }
2304
2305 fn find_remembered_blob(&self, digest: &CacheDigest) -> Option<PathBuf> {
2306 let remembered = self.verified_blobs.lock().unwrap().get(digest).cloned();
2307 if let Some(remembered) = remembered {
2308 if remembered.is_unchanged() {
2319 return Some(remembered.path);
2320 }
2321 self.verified_blobs.lock().unwrap().remove(digest);
2322 }
2323 None
2324 }
2325
2326 fn remember_verified_blob(&self, digest: &CacheDigest, path: &Path) {
2327 let Some(verified) = VerifiedBlob::describe(path) else {
2328 return;
2332 };
2333 self.verified_blobs
2334 .lock()
2335 .unwrap()
2336 .insert(digest.clone(), verified);
2337 }
2338
2339 async fn find_action_result(&self, action: &CacheDigest) -> Result<AgentResponse> {
2340 if let Some(result) = self.actions.find(action)? {
2341 return Ok(AgentResponse::ActionResult {
2342 result: Some(result),
2343 });
2344 }
2345 if !self.remote_mode.reads() {
2346 return Ok(AgentResponse::ActionResult { result: None });
2347 }
2348 let Some(remote) = &self.remote else {
2349 return Ok(AgentResponse::ActionResult { result: None });
2350 };
2351 let lock = self.action_lock(action);
2352 let _guard = lock.lock().await;
2353 if let Some(result) = self.actions.find(action)? {
2354 return Ok(AgentResponse::ActionResult {
2355 result: Some(result),
2356 });
2357 }
2358 if let Some(result) = self
2359 .pending_remote_actions
2360 .lock()
2361 .unwrap()
2362 .get(action)
2363 .cloned()
2364 {
2365 return Ok(AgentResponse::ActionResult {
2366 result: Some(result),
2367 });
2368 }
2369 let _permit = self.remote_transfers.acquire().await?;
2370 match self.get_remote_action_result(remote, action).await {
2371 Ok(Some(result)) => {
2372 self.pending_remote_actions
2373 .lock()
2374 .unwrap()
2375 .insert(action.clone(), result.clone());
2376 Ok(AgentResponse::ActionResult {
2377 result: Some(result),
2378 })
2379 }
2380 Ok(None) => Ok(AgentResponse::ActionResult { result: None }),
2381 Err(error) => {
2382 self.note_remote_failure();
2383 warn!(
2384 "remote cache action lookup failed for {}: {error}",
2385 action.hash
2386 );
2387 Ok(AgentResponse::ActionResult { result: None })
2388 }
2389 }
2390 }
2391
2392 async fn store_action_result(
2393 &self,
2394 result: &RemoteActionResult,
2395 connection: &ConnectionUploads,
2396 ) -> Result<AgentResponse> {
2397 let path = self.actions.store(result)?;
2398 if let Some(uploads) = &self.uploads {
2399 uploads.queue_action_result(result, connection);
2400 }
2401 Ok(AgentResponse::ActionStored { path })
2402 }
2403
2404 async fn join_action_promise(
2405 &self,
2406 adapter: &str,
2407 invocation: &CacheDigest,
2408 ) -> Result<AgentResponse> {
2409 if !self.remote_mode.reads() || !self.remote_mode.writes() {
2410 return Ok(AgentResponse::ActionPromise {
2411 claim: None,
2412 prediction: None,
2413 });
2414 }
2415 let Some(remote) = &self.remote else {
2416 return Ok(AgentResponse::ActionPromise {
2417 claim: None,
2418 prediction: None,
2419 });
2420 };
2421 let deadline = Instant::now() + ACTION_PROMISE_WAIT;
2422 loop {
2423 let _permit = self.remote_transfers.acquire().await?;
2424 let state = remote.join_action_promise(invocation, adapter).await;
2425 drop(_permit);
2426 match state {
2427 Ok(Some(ActionPromiseState::Claimed { claim })) => {
2428 return Ok(AgentResponse::ActionPromise {
2429 claim: Some(claim),
2430 prediction: None,
2431 });
2432 }
2433 Ok(Some(ActionPromiseState::Complete { prediction })) => {
2434 return Ok(AgentResponse::ActionPromise {
2435 claim: None,
2436 prediction: Some(prediction),
2437 });
2438 }
2439 Ok(Some(ActionPromiseState::Pending { retry_after_ms }))
2440 if Instant::now() < deadline =>
2441 {
2442 tokio::time::sleep(Duration::from_millis(retry_after_ms.clamp(10, 5_000)))
2443 .await;
2444 }
2445 Ok(Some(ActionPromiseState::Pending { .. }) | None) => {
2446 return Ok(AgentResponse::ActionPromise {
2447 claim: None,
2448 prediction: None,
2449 });
2450 }
2451 Ok(Some(_)) => {
2452 return Ok(AgentResponse::ActionPromise {
2453 claim: None,
2454 prediction: None,
2455 });
2456 }
2457 Err(error) => {
2458 self.note_remote_failure();
2459 warn!(
2460 "remote cache action promise failed for {}: {error}",
2461 invocation.hash
2462 );
2463 return Ok(AgentResponse::ActionPromise {
2464 claim: None,
2465 prediction: None,
2466 });
2467 }
2468 }
2469 }
2470 }
2471
2472 async fn complete_action_promise(
2473 &self,
2474 claim: &str,
2475 prediction: &ActionPrediction,
2476 ) -> Result<AgentResponse> {
2477 prediction.validate()?;
2478 let Some(remote) = &self.remote else {
2479 return Ok(AgentResponse::ActionPromiseCompleted);
2480 };
2481 let Some(uploads) = &self.uploads else {
2482 return Ok(AgentResponse::ActionPromiseCompleted);
2483 };
2484 if uploads
2485 .wait_for_actions(std::slice::from_ref(&prediction.action))
2486 .await
2487 .contains(&prediction.action)
2488 {
2489 return Ok(AgentResponse::ActionPromiseCompleted);
2492 }
2493 let completion = ActionPromiseCompletion {
2494 claim: claim.to_string(),
2495 prediction: prediction.clone(),
2496 };
2497 let _permit = self.remote_transfers.acquire().await?;
2498 if let Err(error) = remote
2499 .complete_action_promise(&prediction.invocation, &completion)
2500 .await
2501 {
2502 self.note_remote_failure();
2503 warn!(
2504 "remote cache action promise completion failed for {}: {error}",
2505 prediction.invocation.hash
2506 );
2507 }
2508 Ok(AgentResponse::ActionPromiseCompleted)
2509 }
2510
2511 fn note_remote_failure(&self) {
2513 self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
2514 }
2515
2516 async fn get_remote_action_result(
2517 &self,
2518 remote: &RemoteCacheClient,
2519 action: &CacheDigest,
2520 ) -> Result<Option<RemoteActionResult>> {
2521 self.stats
2522 .remote_action_lookups
2523 .fetch_add(1, Ordering::Relaxed);
2524 let _timer = AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
2525 remote.get_action_result(action).await
2526 }
2527
2528 fn record_action_hit(
2529 &self,
2530 action: &CacheDigest,
2531 restore: RestoreStats,
2532 crate_name: Option<String>,
2533 diagnostic: Option<AgentEvent>,
2534 ) -> Result<AgentResponse> {
2535 validate_crate_name(crate_name.as_deref())?;
2536 if self.actions.find(action)?.is_none() {
2537 let pending = self.pending_remote_actions.lock().unwrap().remove(action);
2538 if let Some(result) = pending {
2539 self.actions.store(&result)?;
2540 } else {
2541 bail!("cannot record a hit for a missing action result");
2542 }
2543 }
2544 self.record_restore(restore);
2545 self.stats.lookups.fetch_add(1, Ordering::Relaxed);
2546 self.stats.hits.fetch_add(1, Ordering::Relaxed);
2547 self.emit_action(
2548 diagnostic,
2549 AgentEvent::ActionHit {
2550 crate_name,
2551 restore,
2552 },
2553 );
2554 Ok(AgentResponse::ActionHitRecorded)
2555 }
2556
2557 fn record_restore(&self, restore: RestoreStats) {
2558 self.record_materialization(restore);
2559 atomic_saturating_add(
2560 &self.stats.avoided_compiler_duration_ns,
2561 restore.avoided_compiler_duration_ns,
2562 );
2563 atomic_saturating_add(&self.stats.restored_output_files, restore.output_files);
2564 atomic_saturating_add(&self.stats.restored_output_bytes, restore.output_bytes);
2565 atomic_saturating_add(
2566 &self.stats.reflinked_output_files,
2567 restore.reflinked_output_files,
2568 );
2569 atomic_saturating_add(
2570 &self.stats.reflinked_output_bytes,
2571 restore.reflinked_output_bytes,
2572 );
2573 atomic_saturating_add(
2574 &self.stats.hardlinked_output_files,
2575 restore.hardlinked_output_files,
2576 );
2577 atomic_saturating_add(
2578 &self.stats.hardlinked_output_bytes,
2579 restore.hardlinked_output_bytes,
2580 );
2581 atomic_saturating_add(&self.stats.copied_output_files, restore.copied_output_files);
2582 atomic_saturating_add(&self.stats.copied_output_bytes, restore.copied_output_bytes);
2583 atomic_saturating_add(&self.stats.reused_output_files, restore.reused_output_files);
2584 atomic_saturating_add(&self.stats.reused_output_bytes, restore.reused_output_bytes);
2585 }
2586
2587 fn record_compiler_invocation(
2588 &self,
2589 outcome: &str,
2590 incremental: bool,
2591 crate_name: Option<&str>,
2592 duration_ns: u64,
2593 diagnostic: Option<AgentEvent>,
2594 bypass: Option<String>,
2595 ) -> Result<AgentResponse> {
2596 if !matches!(outcome, "miss" | "unconsulted" | "bypass" | "verification") {
2597 bail!("invalid compiler invocation outcome");
2598 }
2599 validate_crate_name(crate_name)?;
2600 if incremental {
2601 self.stats
2602 .incremental_compilations
2603 .fetch_add(1, Ordering::Relaxed);
2604 }
2605 if matches!(outcome, "miss" | "verification") {
2607 self.stats.lookups.fetch_add(1, Ordering::Relaxed);
2608 }
2609 let mut compiler = self.stats.compiler.lock().unwrap();
2610 let stats = compiler.entry(outcome.to_string()).or_default();
2611 stats.invocations = stats.invocations.saturating_add(1);
2612 stats.duration_ns = stats.duration_ns.saturating_add(duration_ns);
2613 drop(compiler);
2614 if outcome != "verification"
2615 && let Some(crate_name) = crate_name.filter(|name| !name.is_empty())
2616 {
2617 let mut slow = self.stats.slow_compilations.lock().unwrap();
2618 let duration = slow.entry(crate_name.to_string()).or_default();
2619 *duration = duration.saturating_add(duration_ns);
2620 }
2621 let action = match bypass {
2622 Some(kind) => AgentEvent::BypassedCompilation {
2623 kind,
2624 crate_name: crate_name.map(str::to_string),
2625 duration_ns,
2626 },
2627 None => AgentEvent::CompilerInvocation {
2628 outcome: outcome.to_string(),
2629 crate_name: crate_name.map(str::to_string),
2630 duration_ns,
2631 },
2632 };
2633 self.emit_action(diagnostic, action);
2634 Ok(AgentResponse::CompilerInvocationRecorded)
2635 }
2636
2637 fn record_materialization(&self, restore: RestoreStats) {
2638 atomic_saturating_add(&self.stats.materialization_duration_ns, restore.duration_ns);
2639 }
2640
2641 fn record_diagnostic(&self, level: DiagnosticLevel, message: String) -> Result<AgentResponse> {
2652 if message.is_empty()
2653 || message.len() > MAX_DIAGNOSTIC_BYTES
2654 || message.contains(['\n', '\r', '\0'])
2655 {
2656 bail!("invalid shim {}", level.label());
2657 }
2658 let mut diagnostics = self.diagnostics.lock().unwrap();
2659 let key = (level, message.clone());
2660 if !diagnostics.contains(&key) {
2661 if diagnostics.len() >= MAX_DIAGNOSTICS {
2665 bail!("shim diagnostic limit reached");
2666 }
2667 eprintln!("mbx[{}]: {message}", level.label());
2668 self.emit(|| match level {
2669 DiagnosticLevel::Warning => AgentEvent::Warning {
2670 message: message.clone(),
2671 },
2672 DiagnosticLevel::Error => AgentEvent::Error {
2673 message: message.clone(),
2674 },
2675 });
2676 diagnostics.insert(key);
2677 }
2678 Ok(match level {
2679 DiagnosticLevel::Warning => AgentResponse::WarningRecorded,
2680 DiagnosticLevel::Error => AgentResponse::ErrorRecorded,
2681 })
2682 }
2683
2684 fn find_file_digests(
2691 &self,
2692 scope: FileDigestScope,
2693 files: Vec<FileIdentity>,
2694 ) -> Result<AgentResponse> {
2695 if files.len() > MAX_FILE_DIGEST_BATCH {
2696 bail!("too many file-digest lookups in one request");
2697 }
2698 let ledger = self.file_digests.lock().unwrap();
2699 let digests = files
2700 .into_iter()
2701 .map(|file| {
2702 let recorded = ledger.get(&(scope, file.path.clone()))?;
2703 (recorded.file == file).then(|| recorded.digest.clone())
2704 })
2705 .collect();
2706 Ok(AgentResponse::FileDigests { digests })
2707 }
2708
2709 async fn resolve_file_digests(
2712 &self,
2713 scope: FileDigestScope,
2714 files: Vec<FileIdentity>,
2715 ) -> Result<AgentResponse> {
2716 if files.len() > MAX_FILE_DIGEST_BATCH {
2717 bail!("too many file-digest resolutions in one request");
2718 }
2719 for file in &files {
2720 if !file.path.is_absolute() {
2721 bail!("file-digest resolutions need absolute paths");
2722 }
2723 }
2724 let resolutions = stream::iter(
2725 files
2726 .into_iter()
2727 .map(|file| async move { self.resolve_file_digest(scope, file).await }),
2728 )
2729 .buffered(MAX_CONCURRENT_FILE_DIGESTS)
2730 .collect()
2731 .await;
2732 Ok(AgentResponse::FileDigestsResolved { resolutions })
2733 }
2734
2735 async fn resolve_file_digest(
2736 &self,
2737 scope: FileDigestScope,
2738 file: FileIdentity,
2739 ) -> FileDigestResolution {
2740 let lock = {
2741 let key = (scope, file.clone());
2742 let mut locks = self.file_digest_locks.lock().unwrap();
2743 locks.retain(|_, lock| lock.strong_count() > 0);
2744 if let Some(lock) = locks.get(&key).and_then(Weak::upgrade) {
2745 lock
2746 } else {
2747 let lock = Arc::new(FileDigestFlight {
2748 lock: tokio::sync::Mutex::new(()),
2749 resolution: Mutex::new(None),
2750 });
2751 locks.insert(key, Arc::downgrade(&lock));
2752 lock
2753 }
2754 };
2755 let _flight = lock.lock.lock().await;
2756 if let Some(resolution) = lock.resolution.lock().unwrap().clone() {
2757 return resolution;
2758 }
2759 if let Some(digest) = self
2760 .file_digests
2761 .lock()
2762 .unwrap()
2763 .get(&(scope, file.path.clone()))
2764 .filter(|recorded| recorded.file == file)
2765 .map(|recorded| recorded.digest.clone())
2766 .filter(|digest| digest.size == file.len && file.still_describes().unwrap_or(false))
2767 {
2768 return FileDigestResolution::Digest(digest);
2769 }
2770 let Ok(_permit) = self.file_digest_permits.acquire().await else {
2771 return FileDigestResolution::Unresolved;
2772 };
2773 #[cfg(test)]
2774 self.file_digest_reads.fetch_add(1, Ordering::Relaxed);
2775 let path = file.path.clone();
2776 let resolved =
2777 tokio::task::spawn_blocking(move || digest_file_validated(scope, &path)).await;
2778 let Ok(Ok(Some(observed))) = resolved else {
2779 return FileDigestResolution::Unresolved;
2780 };
2781 let reusable = observed.cache_identity.as_ref() == Some(&file);
2785 if !(reusable || observed.racy && observed.identity == file) {
2786 return FileDigestResolution::Unresolved;
2787 }
2788 let resolution = match observed.resolution {
2789 FileDigestResolution::Digest(digest) if digest.size == file.len => {
2790 if reusable {
2791 let _ = self.record_file_digests(
2792 scope,
2793 vec![RecordedFileDigest {
2794 file: file.clone(),
2795 digest: digest.clone(),
2796 }],
2797 );
2798 }
2799 FileDigestResolution::Digest(digest)
2800 }
2801 FileDigestResolution::EmbeddedTimestampMacro => {
2802 FileDigestResolution::EmbeddedTimestampMacro
2803 }
2804 FileDigestResolution::Digest(_) | FileDigestResolution::Unresolved => {
2805 FileDigestResolution::Unresolved
2806 }
2807 };
2808 if reusable {
2809 *lock.resolution.lock().unwrap() = Some(resolution.clone());
2810 }
2811 resolution
2812 }
2813
2814 fn record_file_digests(
2820 &self,
2821 scope: FileDigestScope,
2822 entries: Vec<RecordedFileDigest>,
2823 ) -> Result<AgentResponse> {
2824 if entries.len() > MAX_FILE_DIGEST_BATCH {
2825 bail!("too many file-digest records in one request");
2826 }
2827 for entry in &entries {
2828 validate_file_digest_record(entry)?;
2829 }
2830 let mut ledger = self.file_digests.lock().unwrap();
2831 for entry in entries {
2832 if ledger.len() >= MAX_FILE_DIGEST_ENTRIES
2833 && !ledger.contains_key(&(scope, entry.file.path.clone()))
2834 {
2835 break;
2836 }
2837 ledger.insert((scope, entry.file.path.clone()), entry);
2838 }
2839 Ok(AgentResponse::FileDigestsRecorded)
2840 }
2841
2842 pub fn seed_file_digests(&self, entries: Vec<(FileDigestScope, RecordedFileDigest)>) -> usize {
2851 self.seed_file_digests_with(|| entries)
2852 }
2853
2854 pub fn seed_file_digests_with(
2860 &self,
2861 load: impl FnOnce() -> Vec<(FileDigestScope, RecordedFileDigest)>,
2862 ) -> usize {
2863 let mut ledger = self.file_digests.lock().unwrap();
2864 let entries = load();
2865 let mut seeded = 0;
2866 for (scope, entry) in entries {
2867 if ledger.len() >= MAX_FILE_DIGEST_ENTRIES {
2868 break;
2869 }
2870 if validate_file_digest_record(&entry).is_err() {
2871 continue;
2872 }
2873 if let std::collections::btree_map::Entry::Vacant(slot) =
2874 ledger.entry((scope, entry.file.path.clone()))
2875 {
2876 slot.insert(entry);
2877 seeded += 1;
2878 }
2879 }
2880 seeded
2881 }
2882
2883 pub fn file_digests(&self) -> Vec<(FileDigestScope, RecordedFileDigest)> {
2886 self.file_digests
2887 .lock()
2888 .unwrap()
2889 .iter()
2890 .map(|((scope, _), entry)| (*scope, entry.clone()))
2891 .collect()
2892 }
2893
2894 fn find_action_prediction(
2895 &self,
2896 task: &str,
2897 invocation: &CacheDigest,
2898 ) -> Result<AgentResponse> {
2899 validate_task_identity(task)?;
2900 invocation.validate()?;
2901 let (prediction, prefetch) = {
2902 let mut tasks = self.task_actions.lock().unwrap();
2903 let Some(state) = tasks.get_mut(task) else {
2904 return Ok(AgentResponse::ActionPrediction { prediction: None });
2905 };
2906 activate_prediction_adapter(state, invocation)
2907 };
2908 if let Some(prefetch) = prefetch {
2909 self.spawn_prefetch_predictions(prefetch);
2910 }
2911 Ok(AgentResponse::ActionPrediction { prediction })
2912 }
2913
2914 fn record_action_prediction(
2915 &self,
2916 task: &str,
2917 prediction: ActionPrediction,
2918 ) -> Result<AgentResponse> {
2919 validate_task_identity(task)?;
2920 prediction.validate()?;
2921 let mut tasks = self.task_actions.lock().unwrap();
2922 let state = tasks.entry(task.to_string()).or_default();
2923 if !state
2924 .pending_predictions
2925 .contains_key(&prediction.invocation)
2926 && state.pending_predictions.len() >= MAX_TASK_ACTION_PREDICTIONS
2927 {
2928 bail!("this task run contains too many action predictions");
2929 }
2930 let current = state.predictions.get(&prediction.invocation);
2931 if current != Some(&prediction) {
2932 match state.changed_predictions.get(&prediction.invocation) {
2933 Some(baseline) if baseline.as_ref() == Some(&prediction) => {
2935 state.changed_predictions.remove(&prediction.invocation);
2936 }
2937 Some(_) => {}
2939 None => {
2940 state
2941 .changed_predictions
2942 .insert(prediction.invocation.clone(), current.cloned());
2943 }
2944 }
2945 }
2946 state
2947 .predictions
2948 .insert(prediction.invocation.clone(), prediction.clone());
2949 state
2950 .pending_predictions
2951 .insert(prediction.invocation.clone(), prediction);
2952 Ok(AgentResponse::ActionPredictionRecorded)
2953 }
2954
2955 fn executable_identity_key(
2956 &self,
2957 executable: PathBuf,
2958 environment: BTreeMap<String, Option<String>>,
2959 ) -> Result<ExecutableIdentityKey> {
2960 if !environment.keys().all(|name| {
2967 matches!(
2968 name.as_str(),
2969 "COMPILER_PATH"
2970 | "GCC_EXEC_PREFIX"
2971 | "LIBRARY_PATH"
2972 | "MBX_FUSE_LD"
2973 | "MBX_WORKING_DIRECTORY"
2974 | "PATH"
2975 | "RUSTUP_HOME"
2976 | "RUSTUP_TOOLCHAIN"
2977 | "SDKROOT"
2978 | "MACOSX_DEPLOYMENT_TARGET"
2979 | "LIB"
2980 | "UCRTVersion"
2981 | "UniversalCRTSdkDir"
2982 | "VCToolsInstallDir"
2983 | "VCToolsVersion"
2984 | "WindowsSdkDir"
2985 | "WindowsSDKVersion"
2986 )
2987 }) {
2988 bail!("executable identity contains an unsupported environment variable");
2989 }
2990 Ok(ExecutableIdentityKey {
2991 executable,
2992 environment,
2993 })
2994 }
2995
2996 fn find_executable_identity(
2997 &self,
2998 executable: PathBuf,
2999 environment: BTreeMap<String, Option<String>>,
3000 ) -> Result<AgentResponse> {
3001 let key = self.executable_identity_key(executable, environment)?;
3002 let remembered = self
3003 .executable_identities
3004 .lock()
3005 .unwrap()
3006 .get(&key)
3007 .cloned();
3008 let stdout = match remembered {
3009 Some(stdout) => Some(stdout),
3010 None => {
3011 let persisted = self.load_persisted_identity(&key);
3012 if let Some(stdout) = &persisted {
3013 let _ = self.remember_executable_identity(&key, stdout.clone());
3016 }
3017 persisted
3018 }
3019 };
3020 Ok(AgentResponse::ExecutableIdentity { stdout })
3021 }
3022
3023 fn store_executable_identity(
3024 &self,
3025 executable: PathBuf,
3026 environment: BTreeMap<String, Option<String>>,
3027 stdout: Vec<u8>,
3028 pins: Vec<PinnedFile>,
3029 ) -> Result<AgentResponse> {
3030 if stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE {
3031 bail!("executable identity exceeds {MAX_EXECUTABLE_IDENTITY_SIZE} bytes");
3032 }
3033 let key = self.executable_identity_key(executable, environment)?;
3034 self.remember_executable_identity(&key, stdout.clone())?;
3035 if pins.is_empty() {
3040 } else if !pins.iter().all(PinnedFile::holds) {
3041 debug!("executable identity was not persisted: a pinned file changed during the probe");
3042 } else if let Err(error) = self.persist_executable_identity(&key, pins, &stdout) {
3043 debug!("executable identity was not persisted: {error:#}");
3044 }
3045 Ok(AgentResponse::ExecutableIdentity {
3046 stdout: Some(stdout),
3047 })
3048 }
3049
3050 fn remember_executable_identity(
3051 &self,
3052 key: &ExecutableIdentityKey,
3053 stdout: Vec<u8>,
3054 ) -> Result<()> {
3055 let mut identities = self.executable_identities.lock().unwrap();
3056 let is_new = !identities.contains_key(key);
3057 let previous_size = identities.get(key).map_or(0, Vec::len);
3058 if is_new && identities.len() >= MAX_EXECUTABLE_IDENTITIES {
3059 bail!("executable identity cache contains too many entries");
3060 }
3061 let retained_bytes = identities.values().map(Vec::len).sum::<usize>();
3062 if retained_bytes - previous_size + stdout.len() > MAX_EXECUTABLE_IDENTITY_BYTES {
3063 bail!("executable identity cache contains too many bytes");
3064 }
3065 identities.insert(key.clone(), stdout);
3066 Ok(())
3067 }
3068
3069 fn persisted_identity_path(&self, key: &ExecutableIdentityKey) -> Result<PathBuf> {
3070 let selector = serde_json::to_vec(&(&key.executable, &key.environment))?;
3071 let digest = CacheDigest::blake3(&selector);
3072 Ok(self.identity_dir.join(format!("{}.json", digest.hash)))
3073 }
3074
3075 fn load_persisted_identity(&self, key: &ExecutableIdentityKey) -> Option<Vec<u8>> {
3078 let path = self.persisted_identity_path(key).ok()?;
3079 let bytes = fs::read(path).ok()?;
3080 let persisted: PersistedExecutableIdentity = serde_json::from_slice(&bytes).ok()?;
3081 if persisted.version != PERSISTED_EXECUTABLE_IDENTITY_VERSION
3082 || persisted.executable != key.executable
3083 || persisted.environment != key.environment
3084 || persisted.pins.is_empty()
3085 || persisted.stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE
3086 {
3087 return None;
3088 }
3089 persisted
3090 .pins
3091 .iter()
3092 .all(PinnedFile::holds)
3093 .then(|| persisted.stdout.into_bytes())
3094 }
3095
3096 fn persist_executable_identity(
3097 &self,
3098 key: &ExecutableIdentityKey,
3099 pins: Vec<PinnedFile>,
3100 stdout: &[u8],
3101 ) -> Result<()> {
3102 let Ok(stdout) = std::str::from_utf8(stdout) else {
3104 return Ok(());
3105 };
3106 let persisted = PersistedExecutableIdentity {
3107 version: PERSISTED_EXECUTABLE_IDENTITY_VERSION,
3108 executable: key.executable.clone(),
3109 environment: key.environment.clone(),
3110 pins,
3111 stdout: stdout.to_owned(),
3112 };
3113 let path = self.persisted_identity_path(key)?;
3114 fs::create_dir_all(self.identity_dir.as_path())?;
3115 let mut temporary = tempfile::NamedTempFile::new_in(self.identity_dir.as_path())?;
3117 std::io::Write::write_all(temporary.as_file_mut(), &serde_json::to_vec(&persisted)?)?;
3118 temporary.persist(path).map_err(|error| error.error)?;
3119 Ok(())
3120 }
3121
3122 pub async fn handle_connection<S>(&self, stream: S) -> Result<()>
3124 where
3125 S: AsyncRead + AsyncWrite + Unpin,
3126 {
3127 let (reader, mut writer) = tokio::io::split(stream);
3128 let mut reader = BufReader::new(reader);
3129 let hello = read_request(&mut reader)
3130 .await?
3131 .ok_or_else(|| eyre::eyre!("connection closed before the agent handshake"))?;
3132 let request: AgentRequest = serde_json::from_str(&hello)?;
3133 match request {
3134 AgentRequest::Hello {
3135 protocol,
3136 client_version,
3137 } if protocol == AGENT_PROTOCOL_VERSION && client_version == self.version.as_ref() => {}
3138 AgentRequest::Hello { protocol, .. } if protocol != AGENT_PROTOCOL_VERSION => {
3139 send_response(
3140 &mut writer,
3141 &AgentResponse::Error {
3142 message: format!(
3143 "unsupported agent protocol {protocol}; expected {AGENT_PROTOCOL_VERSION}"
3144 ),
3145 },
3146 )
3147 .await?;
3148 return Ok(());
3149 }
3150 AgentRequest::Hello { client_version, .. } => {
3151 send_response(
3152 &mut writer,
3153 &AgentResponse::Error {
3154 message: format!(
3155 "cache client {client_version} does not match agent {}",
3156 self.version
3157 ),
3158 },
3159 )
3160 .await?;
3161 return Ok(());
3162 }
3163 _ => bail!("the first agent request must be hello"),
3164 }
3165 send_response(
3166 &mut writer,
3167 &AgentResponse::Hello {
3168 protocol: AGENT_PROTOCOL_VERSION,
3169 agent_version: self.version.to_string(),
3170 },
3171 )
3172 .await?;
3173
3174 let mut connection = ConnectionUploads::default();
3178 let served = async {
3179 while let Some(line) = read_request(&mut reader).await? {
3180 let response = match serde_json::from_str(&line) {
3181 Ok(request) => self.respond_on(request, &mut connection).await,
3182 Err(error) => AgentResponse::Error {
3183 message: format!("invalid agent request: {error}"),
3184 },
3185 };
3186 send_response(&mut writer, &response).await?;
3187 }
3188 Ok(())
3189 }
3190 .await;
3191 self.release_bypass(&mut connection);
3193 served
3194 }
3195}
3196
3197async fn read_request<R>(reader: &mut R) -> Result<Option<String>>
3203where
3204 R: AsyncBufRead + Unpin,
3205{
3206 let mut line = Vec::new();
3207 loop {
3208 let available = reader.fill_buf().await?;
3209 if available.is_empty() {
3210 break;
3211 }
3212 let (consumed, complete) = match available.iter().position(|byte| *byte == b'\n') {
3213 Some(index) => (index, true),
3214 None => (available.len(), false),
3215 };
3216 if line.len() + consumed > MAX_REQUEST_BYTES {
3217 bail!("agent request exceeded {MAX_REQUEST_BYTES} bytes");
3218 }
3219 line.extend_from_slice(&available[..consumed]);
3220 reader.consume(consumed + usize::from(complete));
3222 if complete {
3223 return Ok(Some(String::from_utf8(line)?));
3224 }
3225 }
3226 if line.is_empty() {
3227 Ok(None)
3228 } else {
3229 Ok(Some(String::from_utf8(line)?))
3230 }
3231}
3232
3233async fn send_response(
3234 writer: &mut (impl AsyncWrite + Unpin),
3235 response: &AgentResponse,
3236) -> Result<()> {
3237 let mut encoded = serde_json::to_vec(response)?;
3238 encoded.push(b'\n');
3239 writer.write_all(&encoded).await?;
3240 writer.flush().await?;
3241 Ok(())
3242}
3243
3244#[cfg(test)]
3245#[path = "agent_tests.rs"]
3246mod tests;