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, info, 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, digest_file,
37};
38pub use manifest::{is_task_identity, task_manifest_actions};
39use manifest::{
40 merge_remote_task_manifest, merge_task_manifests, task_manifest_dir, validate_task_identity,
41 validate_task_manifest,
42};
43pub use stats::{AgentStats, CompilerStats};
44use wire::MAX_REQUEST_BYTES;
45pub use wire::{
46 AGENT_PROTOCOL_VERSION, ActionDiagnostic, AgentEvent, AgentEventObserver, AgentRequest,
47 AgentResponse, PinnedFile, PinnedState, RestoreStats, WrapperSpan, WrapperTiming,
48};
49
50const MAX_EXECUTABLE_IDENTITIES: usize = 64;
51const MAX_EXECUTABLE_IDENTITY_SIZE: usize = 64 * 1024;
52const MAX_EXECUTABLE_IDENTITY_BYTES: usize = 256 * 1024;
53const TASK_ACTION_MANIFEST_VERSION: u8 = 1;
54const MAX_TASK_ACTION_PREDICTIONS: usize = 16 * 1024;
55const MAX_DIAGNOSTIC_BYTES: usize = 4 * 1024;
57const MAX_DIAGNOSTICS: usize = 128;
63const ACTION_DIAGNOSTIC_PREFIX: &str = "@mbx-action-diagnostic\t";
64
65#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
66enum DiagnosticLevel {
67 Warning,
68 Error,
69}
70
71impl DiagnosticLevel {
72 fn label(self) -> &'static str {
73 match self {
74 Self::Warning => "warning",
75 Self::Error => "error",
76 }
77 }
78}
79
80const MAX_FILE_DIGEST_BATCH: usize = 16 * 1024;
82const MAX_FILE_DIGEST_ENTRIES: usize = 1024 * 1024;
87const MAX_CONCURRENT_FILE_DIGESTS: usize = 8;
90const MAX_REMOTE_TRANSFERS: usize = 64;
91const MAX_PREFETCH_TRANSFERS: usize = 48;
92const MAX_CONCURRENT_BLOB_PACKS: usize = 8;
98const MAX_PREFETCH_ACTIONS: usize = 1024;
109const MAX_PREFETCH_ACTION_WAVE: usize = 32;
114const MAX_PREFETCH_BATCH_LOOKUPS: usize = 1;
120const PREFETCH_ACTION_BATCH_DELAY: Duration = Duration::from_millis(5);
121const MAX_PREFETCH_DIRECTORY_OBJECTS: usize = 100_000;
122const MAX_PREFETCH_OBJECTS_PER_WAVE: usize = 100_000;
123const DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES: u64 = 5 * 1024 * 1024 * 1024;
124
125pub type TaskFallbacks = Arc<dyn Fn() -> Vec<String> + Send + Sync>;
128type FileDigestFlightKey = (FileDigestScope, FileIdentity);
129type FileDigestFlights = BTreeMap<FileDigestFlightKey, Weak<FileDigestFlight>>;
130
131struct FileDigestFlight {
132 lock: tokio::sync::Mutex<()>,
133 resolution: Mutex<Option<FileDigestResolution>>,
134}
135
136const NEWEST_MANIFEST_CANDIDATES: usize = 8;
148
149pub struct AgentRemoteCache {
151 pub client: RemoteCacheClient,
153 pub mode: RemoteCacheMode,
155 pub staging_dir: PathBuf,
157}
158
159#[derive(Default)]
160struct AtomicAgentStats {
161 wrapper_phases_ns: Mutex<BTreeMap<String, u64>>,
162 lookups: AtomicU64,
163 unconsulted: AtomicU64,
164 hits: AtomicU64,
165 stores: AtomicU64,
166 stored_bytes: AtomicU64,
167 verifications: AtomicU64,
168 divergences: AtomicU64,
169 downloaded_bytes: AtomicU64,
170 uploaded_bytes: AtomicU64,
171 background_uploads: AtomicU64,
172 background_upload_failures: AtomicU64,
173 remote_blob_pack_uploads: AtomicU64,
174 remote_blob_pack_upload_blobs: AtomicU64,
175 upload_drain_duration_ns: AtomicU64,
176 prefetched_actions: AtomicU64,
177 predictions_loaded: AtomicU64,
178 remote_failures: AtomicU64,
179 remote_manifest_lookups: AtomicU64,
180 remote_manifest_lookup_duration_ns: AtomicU64,
181 remote_action_lookups: AtomicU64,
182 remote_action_lookup_duration_ns: AtomicU64,
183 remote_blob_requests: AtomicU64,
184 remote_blob_pack_requests: AtomicU64,
185 remote_blob_pack_blobs: AtomicU64,
186 remote_blob_transfer_duration_ns: AtomicU64,
187 local_cas_write_duration_ns: AtomicU64,
188 prefetch_runs: AtomicU64,
189 prefetch_duration_ns: AtomicU64,
190 materialization_duration_ns: AtomicU64,
191 bypasses: Mutex<BTreeMap<String, u64>>,
192 avoided_compiler_duration_ns: AtomicU64,
193 compiler: Mutex<BTreeMap<String, CompilerStats>>,
194 slow_compilations: Mutex<BTreeMap<String, u64>>,
195 restored_output_files: AtomicU64,
196 restored_output_bytes: AtomicU64,
197 reflinked_output_files: AtomicU64,
198 reflinked_output_bytes: AtomicU64,
199 copied_output_files: AtomicU64,
200 copied_output_bytes: AtomicU64,
201 reused_output_files: AtomicU64,
202 reused_output_bytes: AtomicU64,
203}
204
205struct AtomicDurationTimer<'a> {
206 started: Instant,
207 target: &'a AtomicU64,
208}
209
210impl<'a> AtomicDurationTimer<'a> {
211 fn start(target: &'a AtomicU64) -> Self {
212 Self {
213 started: Instant::now(),
214 target,
215 }
216 }
217}
218
219impl Drop for AtomicDurationTimer<'_> {
220 fn drop(&mut self) {
221 atomic_saturating_add(self.target, duration_ns(self.started));
222 }
223}
224
225fn duration_ns(started: Instant) -> u64 {
226 started.elapsed().as_nanos().try_into().unwrap_or(u64::MAX)
227}
228
229fn validate_crate_name(crate_name: Option<&str>) -> Result<()> {
230 if let Some(crate_name) = crate_name
231 && (crate_name.len() > 256 || crate_name.contains(['\0', '\n', '\r']))
232 {
233 bail!("invalid compiler crate name");
234 }
235 Ok(())
236}
237
238#[derive(serde::Deserialize)]
239struct ActionDiagnosticEnvelope {
240 outcome: String,
241 crate_name: Option<String>,
242 diagnostic: ActionDiagnostic,
243}
244
245fn parse_action_diagnostic(
246 message: &str,
247) -> Option<Result<(String, Option<String>, ActionDiagnostic)>> {
248 let payload = message.strip_prefix(ACTION_DIAGNOSTIC_PREFIX)?;
249 Some((|| {
250 let envelope: ActionDiagnosticEnvelope = serde_json::from_str(payload)?;
251 if !matches!(envelope.outcome.as_str(), "hit" | "miss") {
252 bail!("invalid action diagnostic outcome");
253 }
254 validate_crate_name(envelope.crate_name.as_deref())?;
255 Ok((envelope.outcome, envelope.crate_name, envelope.diagnostic))
256 })())
257}
258
259fn atomic_saturating_add(target: &AtomicU64, value: u64) {
260 let _ = target.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
261 Some(current.saturating_add(value))
262 });
263}
264
265fn queue_prefetch_digest(
266 verified: &BTreeMap<CacheDigest, PathBuf>,
267 pending: &mut BTreeMap<CacheDigest, ()>,
268 digest: CacheDigest,
269) {
270 if verified.contains_key(&digest) || pending.contains_key(&digest) {
271 return;
272 }
273 pending.insert(digest, ());
274}
275
276fn queue_prefetch_directory(
277 seen: &BTreeMap<CacheDigest, ()>,
278 pending: &mut BTreeMap<CacheDigest, ()>,
279 digest: CacheDigest,
280 limit: usize,
281) -> bool {
282 if seen.contains_key(&digest) || pending.contains_key(&digest) {
283 return true;
284 }
285 if seen.len().saturating_add(pending.len()) >= limit {
286 return false;
287 }
288 pending.insert(digest, ());
289 true
290}
291
292pub type TaskLoader = dyn for<'a> Fn(&'a CacheAgent, &'a str) -> BoxFuture<'a, ()> + Send + Sync;
294
295#[derive(Clone)]
301pub struct CacheAgent {
302 cas: LocalCas,
303 actions: LocalActionCache,
304 verified_blobs: Arc<Mutex<BTreeMap<CacheDigest, VerifiedBlob>>>,
305 version: Arc<str>,
306 write_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
307 action_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
308 stats: Arc<AtomicAgentStats>,
309 observer: Option<Arc<dyn AgentEventObserver>>,
310 observer_emission: Arc<Mutex<()>>,
311 task_loader: Option<Arc<TaskLoader>>,
314 executable_identities: Arc<Mutex<BTreeMap<ExecutableIdentityKey, Vec<u8>>>>,
315 identity_dir: Arc<PathBuf>,
317 manifest_dir: Arc<PathBuf>,
318 task_actions: Arc<Mutex<BTreeMap<String, TaskActionState>>>,
319 task_fallbacks: Arc<Mutex<BTreeMap<String, TaskFallbacks>>>,
321 next_task_run: Arc<AtomicU64>,
322 manifest_write_lock: Arc<Mutex<()>>,
323 remote: Option<Arc<RemoteCacheClient>>,
324 remote_mode: RemoteCacheMode,
325 remote_staging_dir: Arc<PathBuf>,
326 remote_download_limit: u64,
327 remote_download_bytes: Arc<AtomicU64>,
328 pending_remote_actions: Arc<Mutex<BTreeMap<CacheDigest, RemoteActionResult>>>,
329 remote_transfers: Arc<tokio::sync::Semaphore>,
330 prefetch_transfers: Arc<tokio::sync::Semaphore>,
331 prefetch_tasks: Arc<Mutex<Vec<tokio::task::JoinHandle<()>>>>,
332 diagnostics: Arc<Mutex<BTreeSet<(DiagnosticLevel, String)>>>,
337 file_digests: Arc<Mutex<BTreeMap<(FileDigestScope, PathBuf), RecordedFileDigest>>>,
340 file_digest_locks: Arc<Mutex<FileDigestFlights>>,
342 file_digest_permits: Arc<tokio::sync::Semaphore>,
343 #[cfg(test)]
344 file_digest_reads: Arc<AtomicU64>,
345 uploads: Option<UploadQueue>,
347}
348
349fn validate_file_digest_record(entry: &RecordedFileDigest) -> Result<()> {
352 if !entry.file.path.is_absolute() {
353 bail!("file-digest records need absolute paths");
354 }
355 entry.digest.validate()?;
356 if entry.file.len != entry.digest.size {
357 bail!("file-digest record length does not match its digest");
358 }
359 Ok(())
360}
361
362struct AgentUploadSink {
364 stats: Arc<AtomicAgentStats>,
365}
366
367impl UploadSink for AgentUploadSink {
368 fn record_blob_uploaded(&self, bytes: u64) {
369 self.stats
370 .background_uploads
371 .fetch_add(1, Ordering::Relaxed);
372 self.stats
373 .uploaded_bytes
374 .fetch_add(bytes, Ordering::Relaxed);
375 }
376
377 fn record_action_uploaded(&self) {
378 self.stats
379 .background_uploads
380 .fetch_add(1, Ordering::Relaxed);
381 }
382
383 fn record_blob_pack_uploaded(&self, blobs: u64) {
384 self.stats
385 .remote_blob_pack_uploads
386 .fetch_add(1, Ordering::Relaxed);
387 self.stats
388 .remote_blob_pack_upload_blobs
389 .fetch_add(blobs, Ordering::Relaxed);
390 }
391
392 fn record_upload_failure(&self) {
393 self.stats
394 .background_upload_failures
395 .fetch_add(1, Ordering::Relaxed);
396 self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
397 }
398}
399
400#[derive(Debug, Clone)]
406struct VerifiedBlob {
407 path: PathBuf,
408 len: u64,
409 modified: SystemTime,
410}
411
412impl VerifiedBlob {
413 fn describe(path: &Path) -> Option<Self> {
416 let metadata = std::fs::metadata(path).ok()?;
417 Some(Self {
418 path: path.to_path_buf(),
419 len: metadata.len(),
420 modified: metadata.modified().ok()?,
421 })
422 }
423
424 fn is_unchanged(&self) -> bool {
426 let Ok(metadata) = std::fs::metadata(&self.path) else {
427 return false;
428 };
429 metadata.len() == self.len && metadata.modified().is_ok_and(|now| now == self.modified)
430 }
431}
432
433#[derive(Debug, Clone, Default)]
434struct TaskActionState {
435 manifest: String,
436 baseline_loaded: bool,
437 predictions: BTreeMap<CacheDigest, ActionPrediction>,
438 pending_predictions: BTreeMap<CacheDigest, ActionPrediction>,
439 changed_predictions: BTreeMap<CacheDigest, Option<ActionPrediction>>,
446 prefetched_adapters: BTreeSet<String>,
447 remote_etag: Option<String>,
448}
449
450fn activate_prediction_adapter(
452 state: &mut TaskActionState,
453 invocation: &CacheDigest,
454) -> (Option<ActionPrediction>, Option<Vec<ActionPrediction>>) {
455 let prediction = state.predictions.get(invocation).cloned();
456 let prefetch = prediction.as_ref().and_then(|prediction| {
457 if !state.prefetched_adapters.insert(prediction.adapter.clone()) {
458 return None;
459 }
460 Some(
461 state
462 .predictions
463 .values()
464 .filter(|candidate| candidate.adapter == prediction.adapter)
465 .cloned()
466 .collect(),
467 )
468 });
469 (prediction, prefetch)
470}
471
472struct PrefetchedAction {
473 adapter: String,
474 result: RemoteActionResult,
475}
476
477struct RemoteDownloadReservation {
478 counter: Arc<AtomicU64>,
479 reserved: u64,
480 committed: bool,
481}
482
483impl RemoteDownloadReservation {
484 fn bytes(&self) -> u64 {
485 self.reserved
486 }
487
488 fn commit(mut self, bytes: u64) {
489 debug_assert!(bytes <= self.reserved);
490 self.counter
491 .fetch_sub(self.reserved.saturating_sub(bytes), Ordering::AcqRel);
492 self.committed = true;
493 }
494}
495
496impl Drop for RemoteDownloadReservation {
497 fn drop(&mut self) {
498 if !self.committed {
499 self.counter.fetch_sub(self.reserved, Ordering::AcqRel);
500 }
501 }
502}
503
504#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
505struct ExecutableIdentityKey {
506 executable: PathBuf,
507 environment: BTreeMap<String, Option<String>>,
508}
509
510const PERSISTED_EXECUTABLE_IDENTITY_VERSION: u8 = 2;
512
513#[derive(Debug, Serialize, Deserialize)]
519struct PersistedExecutableIdentity {
520 version: u8,
521 executable: PathBuf,
522 environment: BTreeMap<String, Option<String>>,
523 pins: Vec<PinnedFile>,
524 stdout: String,
525}
526
527impl CacheAgent {
528 pub fn new(cache_dir: impl Into<PathBuf>, version: impl Into<Arc<str>>) -> Self {
530 Self::build(cache_dir.into(), version.into(), None, 0)
531 }
532
533 pub fn new_remote(
535 cache_dir: impl Into<PathBuf>,
536 version: impl Into<Arc<str>>,
537 remote: AgentRemoteCache,
538 ) -> Self {
539 Self::build(
540 cache_dir.into(),
541 version.into(),
542 Some(remote),
543 DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES,
544 )
545 }
546
547 pub fn new_remote_with_download_limit(
549 cache_dir: impl Into<PathBuf>,
550 version: impl Into<Arc<str>>,
551 remote: AgentRemoteCache,
552 max_remote_download_bytes: u64,
553 ) -> Self {
554 Self::build(
555 cache_dir.into(),
556 version.into(),
557 Some(remote),
558 max_remote_download_bytes,
559 )
560 }
561
562 fn build(
563 cache_dir: PathBuf,
564 version: Arc<str>,
565 remote: Option<AgentRemoteCache>,
566 remote_download_limit: u64,
567 ) -> Self {
568 let remote_mode = remote
569 .as_ref()
570 .map_or(RemoteCacheMode::ReadOnly, |remote| remote.mode);
571 let remote_staging_dir = remote.as_ref().map_or_else(
572 || cache_dir.join("remote"),
573 |remote| remote.staging_dir.clone(),
574 );
575 let remote = remote.map(|remote| Arc::new(remote.client));
576 let stats = Arc::new(AtomicAgentStats::default());
577 let remote_transfers = Arc::new(tokio::sync::Semaphore::new(MAX_REMOTE_TRANSFERS));
578 let uploads = remote
581 .clone()
582 .filter(|_| remote_mode.writes())
583 .map(|client| {
584 UploadQueue::new(
585 client,
586 Arc::new(AgentUploadSink {
587 stats: stats.clone(),
588 }),
589 remote_transfers.clone(),
590 )
591 });
592 Self {
593 cas: LocalCas::new(cache_dir.clone()),
594 actions: LocalActionCache::new(cache_dir.clone()),
595 verified_blobs: Arc::new(Mutex::new(BTreeMap::new())),
596 version,
597 write_locks: Arc::new(Mutex::new(BTreeMap::new())),
598 action_locks: Arc::new(Mutex::new(BTreeMap::new())),
599 stats,
600 observer: None,
601 observer_emission: Arc::new(Mutex::new(())),
602 task_loader: None,
603 executable_identities: Arc::new(Mutex::new(BTreeMap::new())),
604 identity_dir: Arc::new(cache_dir.join("executable-identities").join("v1")),
605 manifest_dir: Arc::new(task_manifest_dir(&cache_dir)),
606 task_actions: Arc::new(Mutex::new(BTreeMap::new())),
607 task_fallbacks: Arc::new(Mutex::new(BTreeMap::new())),
608 next_task_run: Arc::new(AtomicU64::new(0)),
609 manifest_write_lock: Arc::new(Mutex::new(())),
610 remote,
611 remote_mode,
612 remote_staging_dir: Arc::new(remote_staging_dir),
613 remote_download_limit,
614 remote_download_bytes: Arc::new(AtomicU64::new(0)),
615 pending_remote_actions: Arc::new(Mutex::new(BTreeMap::new())),
616 remote_transfers,
617 prefetch_transfers: Arc::new(tokio::sync::Semaphore::new(MAX_PREFETCH_TRANSFERS)),
618 prefetch_tasks: Arc::new(Mutex::new(Vec::new())),
619 diagnostics: Arc::new(Mutex::new(BTreeSet::new())),
620 file_digests: Arc::new(Mutex::new(BTreeMap::new())),
621 file_digest_locks: Arc::new(Mutex::new(BTreeMap::new())),
622 file_digest_permits: Arc::new(tokio::sync::Semaphore::new(MAX_CONCURRENT_FILE_DIGESTS)),
623 #[cfg(test)]
624 file_digest_reads: Arc::new(AtomicU64::new(0)),
625 uploads,
626 }
627 }
628
629 #[must_use]
631 pub fn with_observer(mut self, observer: Arc<dyn AgentEventObserver>) -> Self {
632 self.observer = Some(observer);
633 self
634 }
635
636 pub fn with_task_loader(mut self, loader: Arc<TaskLoader>) -> Self {
643 self.task_loader = Some(loader);
644 self
645 }
646
647 async fn ensure_task_loaded(&self, task: &str) {
650 let Some(loader) = &self.task_loader else {
651 return;
652 };
653 if self.task_actions.lock().unwrap().contains_key(task) {
654 return;
655 }
656 loader(self, task).await;
657 }
658
659 fn emit(&self, event: impl FnOnce() -> AgentEvent) {
660 if let Some(observer) = &self.observer {
661 let _emission = self.observer_emission.lock().unwrap();
662 observer.event(event());
663 }
664 }
665
666 fn emit_action(&self, diagnostic: Option<AgentEvent>, action: AgentEvent) {
667 if let Some(observer) = &self.observer {
668 let _emission = self.observer_emission.lock().unwrap();
669 if let Some(diagnostic) = diagnostic {
670 observer.event(diagnostic);
671 }
672 observer.event(action);
673 }
674 }
675
676 fn reserve_remote_download(&self, bytes: u64) -> Result<RemoteDownloadReservation> {
677 let mut current = self.remote_download_bytes.load(Ordering::Acquire);
678 loop {
679 let next = current
680 .checked_add(bytes)
681 .ok_or_else(|| eyre::eyre!("remote cache download budget overflowed"))?;
682 if next > self.remote_download_limit {
683 bail!(
684 "remote cache download budget exceeded: {} bytes requested with {} of {} bytes already used",
685 bytes,
686 current,
687 self.remote_download_limit
688 );
689 }
690 match self.remote_download_bytes.compare_exchange_weak(
691 current,
692 next,
693 Ordering::AcqRel,
694 Ordering::Acquire,
695 ) {
696 Ok(_) => {
697 return Ok(RemoteDownloadReservation {
698 counter: self.remote_download_bytes.clone(),
699 reserved: bytes,
700 committed: false,
701 });
702 }
703 Err(observed) => current = observed,
704 }
705 }
706 }
707
708 fn reserve_remote_download_up_to(&self, requested: u64) -> Result<RemoteDownloadReservation> {
709 let mut current = self.remote_download_bytes.load(Ordering::Acquire);
710 loop {
711 let reserved = requested.min(self.remote_download_limit.saturating_sub(current));
712 let next = current + reserved;
713 match self.remote_download_bytes.compare_exchange_weak(
714 current,
715 next,
716 Ordering::AcqRel,
717 Ordering::Acquire,
718 ) {
719 Ok(_) => {
720 return Ok(RemoteDownloadReservation {
721 counter: self.remote_download_bytes.clone(),
722 reserved,
723 committed: false,
724 });
725 }
726 Err(observed) => current = observed,
727 }
728 }
729 }
730
731 pub async fn begin_task(&self, task: &str) -> Result<String> {
733 self.begin_task_with_remote_errors(task, false, None, true)
734 .await
735 }
736
737 pub async fn begin_task_on_prediction(&self, task: &str) -> Result<String> {
740 self.begin_task_with_remote_errors(task, false, None, false)
741 .await
742 }
743
744 pub async fn begin_session_task(&self, task: &str) -> Result<()> {
750 self.begin_task_with_remote_errors(task, false, Some(task.to_string()), false)
751 .await?;
752 Ok(())
753 }
754
755 pub fn register_task_fallbacks(
766 &self,
767 task: &str,
768 fallbacks: impl Fn() -> Vec<String> + Send + Sync + 'static,
769 ) -> Result<()> {
770 validate_task_identity(task)?;
771 self.task_fallbacks
772 .lock()
773 .unwrap()
774 .insert(task.to_string(), Arc::new(fallbacks));
775 Ok(())
776 }
777
778 pub async fn prefetch_task(&self, task: &str) -> Result<String> {
780 let run = self
781 .begin_task_with_remote_errors(task, true, None, true)
782 .await?;
783 self.wait_for_prefetches().await;
784 Ok(run)
785 }
786
787 async fn begin_task_with_remote_errors(
788 &self,
789 task: &str,
790 strict: bool,
791 run: Option<String>,
792 eager_prefetch: bool,
793 ) -> Result<String> {
794 validate_task_identity(task)?;
795 let early_manifest = {
801 let _write_guard = self.manifest_write_lock.lock().unwrap();
802 let _file_guard = self.lock_task_manifest(task)?;
803 self.load_task_manifest(task)?
804 };
805 let early_actions: BTreeSet<_> = early_manifest
806 .iter()
807 .flat_map(|manifest| manifest.predictions.iter())
808 .map(|prediction| prediction.action.clone())
809 .collect();
810 if eager_prefetch {
811 self.spawn_prefetch_predictions(
812 early_manifest
813 .as_ref()
814 .map(|manifest| manifest.predictions.clone())
815 .unwrap_or_default(),
816 );
817 }
818 let (remote_manifest, mut remote_etag) = if self.remote_mode.reads() {
819 match self.get_remote_task_manifest(task).await {
820 Ok(Some((manifest, etag))) => (Some(manifest), Some(etag)),
821 Ok(None) => (None, None),
822 Err(error) => {
823 if strict {
824 return Err(error).wrap_err_with(|| {
825 format!("remote task action manifest lookup failed for {task}")
826 });
827 }
828 self.note_remote_failure();
829 warn!("remote task action manifest lookup failed for {task}: {error}");
830 (None, None)
831 }
832 }
833 } else {
834 (None, None)
835 };
836 let manifest = if self.remote.is_some() && self.remote_mode.reads() {
840 let _write_guard = self.manifest_write_lock.lock().unwrap();
841 let _file_guard = self.lock_task_manifest(task)?;
842 let local_manifest = self.load_task_manifest(task)?;
843 let manifest = match (remote_manifest, local_manifest) {
844 (Some(remote), Some(local)) => {
845 let (manifest, merged) = merge_remote_task_manifest(task, remote, local);
846 if !merged {
847 remote_etag = None;
848 }
849 Some(manifest)
850 }
851 (Some(remote), None) => Some(remote),
852 (None, local) => local,
853 };
854 if let Some(manifest) = &manifest {
855 self.persist_task_manifest(manifest)?;
856 }
857 manifest
858 } else {
859 early_manifest
860 };
861 let (manifest, remote_etag) = match manifest {
862 Some(manifest) => (Some(manifest), remote_etag),
863 None => (self.inherit_task_manifest(task, strict).await?, None),
866 };
867 let mut state = if let Some(manifest) = manifest {
868 TaskActionState {
869 manifest: task.to_string(),
870 baseline_loaded: true,
871 predictions: manifest
872 .predictions
873 .into_iter()
874 .map(|prediction| (prediction.invocation.clone(), prediction))
875 .collect(),
876 pending_predictions: BTreeMap::new(),
877 changed_predictions: BTreeMap::new(),
878 prefetched_adapters: BTreeSet::new(),
879 remote_etag,
880 }
881 } else {
882 TaskActionState {
883 manifest: task.to_string(),
884 baseline_loaded: true,
885 remote_etag,
886 ..TaskActionState::default()
887 }
888 };
889 let run = run.unwrap_or_else(|| {
890 let sequence = self.next_task_run.fetch_add(1, Ordering::Relaxed);
891 CacheDigest::blake3(format!("{task}\0{}\0{sequence}", std::process::id()).as_bytes())
892 .hash
893 });
894 self.stats.predictions_loaded.fetch_max(
898 state.predictions.len().try_into().unwrap_or(u64::MAX),
899 Ordering::Relaxed,
900 );
901 if eager_prefetch {
902 state.prefetched_adapters.extend(
903 state
904 .predictions
905 .values()
906 .map(|prediction| prediction.adapter.clone()),
907 );
908 }
909 let predictions = state
913 .predictions
914 .values()
915 .filter(|prediction| !early_actions.contains(&prediction.action))
916 .cloned()
917 .collect();
918 self.task_actions.lock().unwrap().insert(run.clone(), state);
919 if eager_prefetch {
920 self.spawn_prefetch_predictions(predictions);
921 }
922 Ok(run)
923 }
924
925 async fn inherit_task_manifest(
933 &self,
934 task: &str,
935 strict: bool,
936 ) -> Result<Option<TaskActionManifest>> {
937 let fallbacks = self.task_fallbacks.lock().unwrap().get(task).cloned();
938 let Some(fallbacks) = fallbacks else {
939 return Ok(None);
940 };
941 let named = tokio::task::spawn_blocking(move || fallbacks()).await?;
942 let mut tried = BTreeSet::from([task.to_string()]);
943 let candidates = named.into_iter().map(|identity| (identity, true)).chain(
944 self.newest_task_identities()
945 .into_iter()
946 .map(|identity| (identity, false)),
947 );
948 for (identity, named) in candidates {
949 if !tried.insert(identity.clone()) || validate_task_identity(&identity).is_err() {
950 continue;
951 }
952 let mut found = self.load_task_manifest(&identity)?;
953 if found.is_none() && named && self.remote_mode.reads() {
954 found = match self.get_remote_task_manifest(&identity).await {
955 Ok(manifest) => manifest.map(|(manifest, _)| manifest),
956 Err(error) => {
957 if strict {
958 return Err(error).wrap_err_with(|| {
959 format!("remote task action manifest lookup failed for {identity}")
960 });
961 }
962 self.note_remote_failure();
963 warn!("remote task action manifest lookup failed for {identity}: {error}");
964 None
965 }
966 };
967 }
968 let Some(mut manifest) = found else {
969 continue;
970 };
971 if manifest.predictions.is_empty()
972 || (!named && manifest.predictions.len() > MAX_TASK_ACTION_PREDICTIONS / 2)
973 {
974 continue;
975 }
976 info!(
977 "no predictions were recorded for {task}; inheriting {} from {identity}",
978 manifest.predictions.len()
979 );
980 manifest.task = task.to_string();
981 validate_task_manifest(&manifest, task)?;
982 let _write_guard = self.manifest_write_lock.lock().unwrap();
983 let _file_guard = self.lock_task_manifest(task)?;
984 if let Some(recorded) = self.load_task_manifest(task)? {
988 return Ok(Some(recorded));
989 }
990 self.persist_task_manifest(&manifest)?;
991 return Ok(Some(manifest));
992 }
993 Ok(None)
994 }
995
996 fn newest_task_identities(&self) -> Vec<String> {
998 let Ok(entries) = fs::read_dir(self.manifest_dir.as_path()) else {
999 return Vec::new();
1000 };
1001 let mut manifests: Vec<(std::time::SystemTime, String)> = entries
1002 .filter_map(|entry| {
1003 let entry = entry.ok()?;
1004 let name = entry.file_name();
1005 let identity = name.to_str()?.strip_suffix(".json")?.to_string();
1006 validate_task_identity(&identity).ok()?;
1007 let modified = entry.metadata().ok()?.modified().ok()?;
1008 Some((modified, identity))
1009 })
1010 .collect();
1011 manifests.sort_by(|left, right| right.cmp(left));
1012 manifests
1013 .into_iter()
1014 .take(NEWEST_MANIFEST_CANDIDATES)
1015 .map(|(_, identity)| identity)
1016 .collect()
1017 }
1018
1019 pub async fn cancel_prefetches(&self) {
1021 let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
1022 for task in &tasks {
1023 task.abort();
1024 }
1025 for task in tasks {
1026 if let Err(error) = task.await
1027 && !error.is_cancelled()
1028 {
1029 warn!("remote action prefetch task failed: {error}");
1030 }
1031 }
1032 }
1033
1034 pub async fn wait_for_uploads(&self) {
1041 let Some(uploads) = &self.uploads else {
1042 return;
1043 };
1044 let _timer = AtomicDurationTimer::start(&self.stats.upload_drain_duration_ns);
1045 uploads.drain().await;
1046 }
1047
1048 async fn wait_for_prefetches(&self) {
1049 let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
1050 for task in tasks {
1051 if let Err(error) = task.await {
1052 warn!("remote action prefetch task failed: {error}");
1053 }
1054 }
1055 }
1056
1057 pub async fn commit_task(&self, run: &str) -> Result<()> {
1059 self.commit_task_actions(run).await.map(|_| ())
1060 }
1061
1062 pub async fn commit_task_actions(&self, run: &str) -> Result<Vec<ActionPrediction>> {
1068 validate_task_identity(run)?;
1069 let state = {
1070 let mut runs = self.task_actions.lock().unwrap();
1071 let state = runs
1072 .get(run)
1073 .ok_or_else(|| eyre::eyre!("task action manifest baseline was not loaded"))?;
1074 if !state.baseline_loaded {
1075 bail!("task action manifest baseline was not loaded");
1076 }
1077 debug!(
1083 "committing {} changed of {} recorded predictions for {}",
1084 state.changed_predictions.len(),
1085 state.pending_predictions.len(),
1086 state.manifest
1087 );
1088 if state.changed_predictions.is_empty()
1089 && (self.remote.is_none() || !self.remote_mode.writes())
1090 {
1091 let completed = state.pending_predictions.values().cloned().collect();
1092 runs.remove(run);
1093 return Ok(completed);
1094 }
1095 state.clone()
1096 };
1097 let task = state.manifest;
1098 validate_task_identity(&task)?;
1099 let completed = state
1100 .pending_predictions
1101 .values()
1102 .cloned()
1103 .collect::<Vec<_>>();
1104 let (manifest, introduced) = {
1105 let _write_guard = self.manifest_write_lock.lock().unwrap();
1106 let _file_guard = self.lock_task_manifest(&task)?;
1107 let mut predictions = self
1108 .load_task_manifest(&task)?
1109 .map(|manifest| {
1110 manifest
1111 .predictions
1112 .into_iter()
1113 .map(|prediction| (prediction.invocation.clone(), prediction))
1114 .collect::<BTreeMap<_, _>>()
1115 })
1116 .unwrap_or_default();
1117 let introduced: BTreeSet<CacheDigest> = state
1122 .pending_predictions
1123 .keys()
1124 .filter(|invocation| !predictions.contains_key(*invocation))
1125 .cloned()
1126 .collect();
1127 predictions.extend(state.pending_predictions);
1132 let manifest = TaskActionManifest {
1133 version: TASK_ACTION_MANIFEST_VERSION,
1134 task: task.clone(),
1135 predictions: predictions.into_values().collect(),
1136 };
1137 validate_task_manifest(&manifest, &task)?;
1138 self.persist_task_manifest(&manifest)?;
1139 (manifest, introduced)
1140 };
1141 self.task_actions.lock().unwrap().remove(run);
1142 if self.remote_mode.writes() {
1143 let mut manifest = manifest;
1147 if let Some(uploads) = &self.uploads {
1148 let actions: Vec<CacheDigest> = manifest
1149 .predictions
1150 .iter()
1151 .map(|prediction| prediction.action.clone())
1152 .collect();
1153 let unpublished = uploads.wait_for_actions(&actions).await;
1154 let withheld = manifest
1155 .predictions
1156 .iter()
1157 .filter(|prediction| {
1158 introduced.contains(&prediction.invocation)
1159 && unpublished.contains(&prediction.action)
1160 })
1161 .count();
1162 if withheld > 0 {
1163 warn!(
1166 "{withheld} of {} predicted actions were not published, so the remote task action manifest omits them",
1167 manifest.predictions.len()
1168 );
1169 manifest.predictions.retain(|prediction| {
1170 !(introduced.contains(&prediction.invocation)
1171 && unpublished.contains(&prediction.action))
1172 });
1173 }
1174 }
1175 match self
1176 .put_remote_task_manifest(&task, manifest, state.remote_etag)
1177 .await
1178 {
1179 Ok(remote_manifest) => {
1180 let _write_guard = self.manifest_write_lock.lock().unwrap();
1181 let reconciliation = (|| {
1182 let _file_guard = self.lock_task_manifest(&task)?;
1183 let manifest = match self.load_task_manifest(&task)? {
1184 Some(local) => {
1185 merge_remote_task_manifest(&task, remote_manifest, local).0
1186 }
1187 None => remote_manifest,
1188 };
1189 self.persist_task_manifest(&manifest)
1190 })();
1191 if let Err(error) = reconciliation {
1192 warn!(
1193 "remote task action manifest reconciliation failed for {task}: {error}"
1194 );
1195 }
1196 }
1197 Err(error) => {
1198 self.note_remote_failure();
1199 warn!("remote task action manifest upload failed for {task}: {error}");
1200 }
1201 }
1202 }
1203 Ok(completed)
1204 }
1205
1206 fn task_manifest_path(&self, task: &str) -> PathBuf {
1207 self.manifest_dir.join(format!("{task}.json"))
1208 }
1209
1210 fn task_manifest_lock_path(&self, task: &str) -> PathBuf {
1211 self.manifest_dir.join("locks").join(format!("{task}.lock"))
1212 }
1213
1214 fn lock_task_manifest(&self, task: &str) -> Result<fslock::LockFile> {
1215 let path = self.task_manifest_lock_path(task);
1216 fs::create_dir_all(path.parent().expect("task manifest lock has a parent"))?;
1217 let mut lock = fslock::LockFile::open(&path)?;
1218 lock.lock()?;
1219 Ok(lock)
1220 }
1221
1222 fn load_task_manifest(&self, task: &str) -> Result<Option<TaskActionManifest>> {
1223 match fs::read(self.task_manifest_path(task)) {
1224 Ok(contents) => Ok(Some(self.parse_task_manifest(task, &contents, false)?)),
1225 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
1226 Err(error) => Err(error.into()),
1227 }
1228 }
1229
1230 fn parse_task_manifest(
1231 &self,
1232 task: &str,
1233 contents: &[u8],
1234 require_canonical: bool,
1235 ) -> Result<TaskActionManifest> {
1236 let manifest: TaskActionManifest = serde_json::from_slice(contents)?;
1237 validate_task_manifest(&manifest, task)?;
1238 if require_canonical && canonical_json(&manifest)? != contents {
1239 bail!("task action manifest is not canonical JSON");
1240 }
1241 Ok(manifest)
1242 }
1243
1244 fn task_manifest_selector(task: &str) -> Result<(Vec<u8>, CacheDigest)> {
1245 TaskActionManifest::selector(task)
1246 }
1247
1248 fn persist_task_manifest(&self, manifest: &TaskActionManifest) -> Result<()> {
1249 let bytes = canonical_json(manifest)?;
1250 fs::create_dir_all(self.manifest_dir.as_path())?;
1251 let mut temporary = tempfile::NamedTempFile::new_in(self.manifest_dir.as_path())?;
1252 std::io::Write::write_all(temporary.as_file_mut(), &bytes)?;
1253 temporary.as_file_mut().sync_all()?;
1254 temporary
1255 .persist(self.task_manifest_path(&manifest.task))
1256 .map_err(|error| error.error)?;
1257 Ok(())
1258 }
1259
1260 async fn get_remote_task_manifest(
1261 &self,
1262 task: &str,
1263 ) -> Result<Option<(TaskActionManifest, String)>> {
1264 let Some(remote) = &self.remote else {
1265 return Ok(None);
1266 };
1267 let (_, selector) = Self::task_manifest_selector(task)?;
1268 let _permit = self.remote_transfers.acquire().await?;
1269 self.stats
1270 .remote_manifest_lookups
1271 .fetch_add(1, Ordering::Relaxed);
1272 let _timer = AtomicDurationTimer::start(&self.stats.remote_manifest_lookup_duration_ns);
1273 let Some(remote_manifest) = remote.get_action_manifest(&selector).await? else {
1274 return Ok(None);
1275 };
1276 let manifest = self.parse_task_manifest(task, &remote_manifest.bytes, true)?;
1277 Ok(Some((manifest, remote_manifest.etag)))
1278 }
1279
1280 async fn put_remote_task_manifest(
1281 &self,
1282 task: &str,
1283 mut manifest: TaskActionManifest,
1284 mut expected_etag: Option<String>,
1285 ) -> Result<TaskActionManifest> {
1286 let Some(remote) = &self.remote else {
1287 return Ok(manifest);
1288 };
1289 let (_, selector) = Self::task_manifest_selector(task)?;
1290 for _ in 0..4 {
1291 let bytes = canonical_json(&manifest)?;
1292 let outcome = {
1293 let _permit = self.remote_transfers.acquire().await?;
1294 remote
1295 .put_action_manifest(&selector, &bytes, expected_etag.as_deref())
1296 .await?
1297 };
1298 match outcome {
1299 ManifestPutOutcome::Stored => return Ok(manifest),
1300 ManifestPutOutcome::PreconditionFailed => {
1301 let Some((remote_manifest, etag)) = self.get_remote_task_manifest(task).await?
1302 else {
1303 expected_etag = None;
1304 continue;
1305 };
1306 manifest = merge_task_manifests(task, Some(remote_manifest), manifest)?;
1307 expected_etag = Some(etag);
1308 }
1309 }
1310 }
1311 bail!("remote task action manifest changed too frequently")
1312 }
1313
1314 pub fn stats(&self) -> AgentStats {
1316 AgentStats {
1317 session_duration_ns: 0,
1318 lookups: self.stats.lookups.load(Ordering::Relaxed),
1319 unconsulted: self.stats.unconsulted.load(Ordering::Relaxed),
1320 hits: self.stats.hits.load(Ordering::Relaxed),
1321 stores: self.stats.stores.load(Ordering::Relaxed),
1322 stored_bytes: self.stats.stored_bytes.load(Ordering::Relaxed),
1323 verifications: self.stats.verifications.load(Ordering::Relaxed),
1324 divergences: self.stats.divergences.load(Ordering::Relaxed),
1325 downloaded_bytes: self.stats.downloaded_bytes.load(Ordering::Relaxed),
1326 uploaded_bytes: self.stats.uploaded_bytes.load(Ordering::Relaxed),
1327 background_uploads: self.stats.background_uploads.load(Ordering::Relaxed),
1328 background_upload_failures: self
1329 .stats
1330 .background_upload_failures
1331 .load(Ordering::Relaxed),
1332 remote_blob_pack_uploads: self.stats.remote_blob_pack_uploads.load(Ordering::Relaxed),
1333 remote_blob_pack_upload_blobs: self
1334 .stats
1335 .remote_blob_pack_upload_blobs
1336 .load(Ordering::Relaxed),
1337 upload_drain_duration_ns: self.stats.upload_drain_duration_ns.load(Ordering::Relaxed),
1338 prefetched_actions: self.stats.prefetched_actions.load(Ordering::Relaxed),
1339 predictions_loaded: self.stats.predictions_loaded.load(Ordering::Relaxed),
1340 bypasses: self.stats.bypasses.lock().unwrap().clone(),
1341 avoided_compiler_duration_ns: self
1342 .stats
1343 .avoided_compiler_duration_ns
1344 .load(Ordering::Relaxed),
1345 wrapper_phases_ns: self.stats.wrapper_phases_ns.lock().unwrap().clone(),
1346 compiler: self.stats.compiler.lock().unwrap().clone(),
1347 slow_compilations: self.stats.slow_compilations.lock().unwrap().clone(),
1348 remote_failures: self.stats.remote_failures.load(Ordering::Relaxed),
1349 remote_manifest_lookups: self.stats.remote_manifest_lookups.load(Ordering::Relaxed),
1350 remote_manifest_lookup_duration_ns: self
1351 .stats
1352 .remote_manifest_lookup_duration_ns
1353 .load(Ordering::Relaxed),
1354 remote_action_lookups: self.stats.remote_action_lookups.load(Ordering::Relaxed),
1355 remote_action_lookup_duration_ns: self
1356 .stats
1357 .remote_action_lookup_duration_ns
1358 .load(Ordering::Relaxed),
1359 remote_blob_requests: self.stats.remote_blob_requests.load(Ordering::Relaxed),
1360 remote_blob_pack_requests: self.stats.remote_blob_pack_requests.load(Ordering::Relaxed),
1361 remote_blob_pack_blobs: self.stats.remote_blob_pack_blobs.load(Ordering::Relaxed),
1362 remote_blob_transfer_duration_ns: self
1363 .stats
1364 .remote_blob_transfer_duration_ns
1365 .load(Ordering::Relaxed),
1366 local_cas_write_duration_ns: self
1367 .stats
1368 .local_cas_write_duration_ns
1369 .load(Ordering::Relaxed),
1370 prefetch_runs: self.stats.prefetch_runs.load(Ordering::Relaxed),
1371 prefetch_duration_ns: self.stats.prefetch_duration_ns.load(Ordering::Relaxed),
1372 materialization_duration_ns: self
1373 .stats
1374 .materialization_duration_ns
1375 .load(Ordering::Relaxed),
1376 restored_output_files: self.stats.restored_output_files.load(Ordering::Relaxed),
1377 restored_output_bytes: self.stats.restored_output_bytes.load(Ordering::Relaxed),
1378 reflinked_output_files: self.stats.reflinked_output_files.load(Ordering::Relaxed),
1379 reflinked_output_bytes: self.stats.reflinked_output_bytes.load(Ordering::Relaxed),
1380 copied_output_files: self.stats.copied_output_files.load(Ordering::Relaxed),
1381 copied_output_bytes: self.stats.copied_output_bytes.load(Ordering::Relaxed),
1382 reused_output_files: self.stats.reused_output_files.load(Ordering::Relaxed),
1383 reused_output_bytes: self.stats.reused_output_bytes.load(Ordering::Relaxed),
1384 }
1385 }
1386
1387 pub async fn handle_requests(
1394 &self,
1395 requests: impl IntoIterator<Item = AgentRequest>,
1396 ) -> Vec<AgentResponse> {
1397 let mut connection = ConnectionUploads::default();
1398 let mut responses = Vec::new();
1399 for request in requests {
1400 responses.push(self.respond_on(request, &mut connection).await);
1401 }
1402 responses
1403 }
1404
1405 fn write_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1406 Self::digest_lock(&self.write_locks, digest)
1407 }
1408
1409 fn action_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1410 Self::digest_lock(&self.action_locks, digest)
1411 }
1412
1413 fn digest_lock(
1414 locks: &Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>,
1415 digest: &CacheDigest,
1416 ) -> Arc<tokio::sync::Mutex<()>> {
1417 let mut locks = locks.lock().unwrap();
1418 locks.retain(|_, lock| lock.strong_count() > 0);
1419 if let Some(lock) = locks.get(digest).and_then(Weak::upgrade) {
1420 return lock;
1421 }
1422 let lock = Arc::new(tokio::sync::Mutex::new(()));
1423 locks.insert(digest.clone(), Arc::downgrade(&lock));
1424 lock
1425 }
1426
1427 #[cfg(test)]
1432 async fn respond(&self, request: AgentRequest) -> AgentResponse {
1433 self.respond_on(request, &mut ConnectionUploads::default())
1434 .await
1435 }
1436
1437 async fn respond_on(
1438 &self,
1439 request: AgentRequest,
1440 connection: &mut ConnectionUploads,
1441 ) -> AgentResponse {
1442 match &request {
1443 AgentRequest::FindActionPrediction { task, .. }
1444 | AgentRequest::RecordActionPrediction { task, .. } => {
1445 self.ensure_task_loaded(task).await;
1446 }
1447 _ => {}
1448 }
1449 let result = match request {
1450 AgentRequest::BeginTask { task } => self
1451 .begin_task(&task)
1452 .await
1453 .map(|run| AgentResponse::TaskBegun { run }),
1454 AgentRequest::CommitTask { run } => self
1455 .commit_task(&run)
1456 .await
1457 .map(|()| AgentResponse::TaskCommitted),
1458 AgentRequest::FindBlob { digest } => self.find_blob(&digest).await,
1459 AgentRequest::FindBlobs { digests } => self.find_blobs(digests).await,
1460 AgentRequest::StoreBlob { digest, source } => {
1461 self.store_blob(&digest, &source, connection).await
1462 }
1463 AgentRequest::FindActionResult { action } => {
1464 self.stats.lookups.fetch_add(1, Ordering::Relaxed);
1465 self.find_action_result(&action).await
1466 }
1467 AgentRequest::RecordActionHit {
1468 action,
1469 restore,
1470 crate_name,
1471 } => {
1472 let diagnostic = connection
1473 .take_action_diagnostic("hit", crate_name.as_deref())
1474 .map(|diagnostic| AgentEvent::ActionDiagnostic {
1475 outcome: "hit".into(),
1476 crate_name: crate_name.clone(),
1477 diagnostic,
1478 });
1479 self.record_action_hit(&action, restore, crate_name, diagnostic)
1480 }
1481 AgentRequest::RecordBypass { kind } => {
1482 *self
1483 .stats
1484 .bypasses
1485 .lock()
1486 .unwrap()
1487 .entry(kind.clone())
1488 .or_insert(0) += 1;
1489 self.emit(|| AgentEvent::Bypass { kind });
1490 Ok(AgentResponse::BypassRecorded)
1491 }
1492 AgentRequest::RecordUnconsulted => {
1493 self.stats.unconsulted.fetch_add(1, Ordering::Relaxed);
1494 self.emit(|| AgentEvent::Unconsulted);
1495 Ok(AgentResponse::UnconsultedRecorded)
1496 }
1497 AgentRequest::RecordWarning { message } => match parse_action_diagnostic(&message) {
1498 Some(Ok((outcome, crate_name, diagnostic))) => {
1499 connection.record_action_diagnostic(outcome, crate_name, diagnostic);
1500 Ok(AgentResponse::WarningRecorded)
1501 }
1502 Some(Err(error)) => Err(error),
1503 None => self.record_diagnostic(DiagnosticLevel::Warning, message),
1504 },
1505 AgentRequest::RecordError { message } => {
1506 self.record_diagnostic(DiagnosticLevel::Error, message)
1507 }
1508 AgentRequest::FindFileDigests { scope, files } => self.find_file_digests(scope, files),
1509 AgentRequest::JoinActionPromise {
1510 adapter,
1511 invocation,
1512 } => self.join_action_promise(&adapter, &invocation).await,
1513 AgentRequest::CompleteActionPromise { claim, prediction } => {
1514 self.complete_action_promise(&claim, &prediction).await
1515 }
1516 AgentRequest::ResolveFileDigests { scope, files } => {
1517 self.resolve_file_digests(scope, files).await
1518 }
1519 AgentRequest::RecordFileDigests { scope, entries } => {
1520 self.record_file_digests(scope, entries)
1521 }
1522 AgentRequest::RecordWrapperTiming { timing } => (|| {
1523 if timing.spans.len() > 512
1524 || timing.phases_ns.len() > 32
1525 || timing.adapter.len() > 64
1526 || timing.unit.as_ref().is_some_and(|s| s.len() > 1024)
1527 || timing.phases_ns.keys().any(|s| s.len() > 64)
1528 || timing.spans.iter().any(|s| s.name.len() > 64)
1529 {
1530 bail!("wrapper timing exceeds telemetry limits");
1531 }
1532 {
1533 let mut phases = self.stats.wrapper_phases_ns.lock().unwrap();
1534 for (phase, duration) in &timing.phases_ns {
1535 let total = phases.entry(phase.clone()).or_default();
1536 *total = total.saturating_add(*duration);
1537 }
1538 }
1539 self.emit(|| AgentEvent::WrapperTiming { timing });
1540 Ok(AgentResponse::WrapperTimingRecorded)
1541 })(),
1542 AgentRequest::RecordCompilerInvocation {
1543 outcome,
1544 crate_name,
1545 duration_ns,
1546 } => {
1547 let diagnostic = connection
1548 .take_action_diagnostic(&outcome, crate_name.as_deref())
1549 .map(|diagnostic| AgentEvent::ActionDiagnostic {
1550 outcome: outcome.clone(),
1551 crate_name: crate_name.clone(),
1552 diagnostic,
1553 });
1554 self.record_compiler_invocation(
1555 &outcome,
1556 crate_name.as_deref(),
1557 duration_ns,
1558 diagnostic,
1559 )
1560 }
1561 AgentRequest::RecordActionVerification { matched, restore } => {
1562 self.record_materialization(restore);
1563 self.stats.verifications.fetch_add(1, Ordering::Relaxed);
1564 if !matched {
1565 self.stats.divergences.fetch_add(1, Ordering::Relaxed);
1566 }
1567 self.emit(|| AgentEvent::Verification { matched, restore });
1568 Ok(AgentResponse::ActionVerificationRecorded)
1569 }
1570 AgentRequest::StoreActionResult { result } => {
1571 self.store_action_result(&result, connection).await
1572 }
1573 AgentRequest::FindActionPrediction { task, invocation } => {
1574 self.find_action_prediction(&task, &invocation)
1575 }
1576 AgentRequest::RecordActionPrediction { task, prediction } => {
1577 self.record_action_prediction(&task, prediction)
1578 }
1579 AgentRequest::FindExecutableIdentity {
1580 executable,
1581 environment,
1582 } => self.find_executable_identity(executable, environment),
1583 AgentRequest::StoreExecutableIdentity {
1584 executable,
1585 environment,
1586 stdout,
1587 pins,
1588 } => self.store_executable_identity(executable, environment, stdout, pins),
1589 AgentRequest::Hello { .. } => {
1590 Err(eyre::eyre!("hello is only valid as the first request"))
1591 }
1592 };
1593 result.unwrap_or_else(|error| AgentResponse::Error {
1594 message: error.to_string(),
1595 })
1596 }
1597
1598 async fn find_blob(&self, digest: &CacheDigest) -> Result<AgentResponse> {
1599 if let Some(path) = self.find_verified_blob(digest)? {
1600 return Ok(AgentResponse::Blob { path: Some(path) });
1601 }
1602 if !self.remote_mode.reads() {
1603 return Ok(AgentResponse::Blob { path: None });
1604 }
1605 let Some(remote) = &self.remote else {
1606 return Ok(AgentResponse::Blob { path: None });
1607 };
1608 match self.fetch_remote_blob(remote, digest).await {
1609 Ok(path) => Ok(AgentResponse::Blob { path: Some(path) }),
1610 Err(error) => {
1611 warn!(
1612 "remote cache blob lookup failed for {}: {error}",
1613 digest.hash
1614 );
1615 Ok(AgentResponse::Blob { path: None })
1616 }
1617 }
1618 }
1619
1620 async fn find_blobs(&self, digests: Vec<CacheDigest>) -> Result<AgentResponse> {
1621 let mut paths = BTreeMap::new();
1622 let mut missing = Vec::new();
1623 for digest in &digests {
1624 match self.find_verified_blob(digest)? {
1625 Some(path) => {
1626 paths.insert(digest.clone(), path);
1627 }
1628 None => {
1629 missing.push(digest.clone());
1630 }
1631 }
1632 }
1633
1634 if !missing.is_empty()
1635 && self.remote_mode.reads()
1636 && let Some(remote) = &self.remote
1637 {
1638 paths.extend(self.fetch_remote_blobs(remote, missing, None).await);
1639 }
1640
1641 Ok(AgentResponse::Blobs {
1642 paths: digests
1643 .into_iter()
1644 .map(|digest| paths.get(&digest).cloned())
1645 .collect(),
1646 })
1647 }
1648
1649 async fn store_blob(
1650 &self,
1651 digest: &CacheDigest,
1652 source: &Path,
1653 connection: &mut ConnectionUploads,
1654 ) -> Result<AgentResponse> {
1655 let path = {
1656 let lock = self.write_lock(digest);
1657 let _guard = lock.lock().await;
1658 if let Some(path) = self.find_verified_blob(digest)? {
1659 path
1660 } else {
1661 let path = self.cas.store_file(digest, source)?;
1662 self.remember_verified_blob(digest, &path);
1663 self.stats.stores.fetch_add(1, Ordering::Relaxed);
1664 self.stats
1665 .stored_bytes
1666 .fetch_add(digest.size, Ordering::Relaxed);
1667 path
1668 }
1669 };
1670 if let Some(uploads) = &self.uploads {
1673 uploads.queue_blob(digest, path.clone(), connection);
1674 }
1675 Ok(AgentResponse::Stored { path })
1676 }
1677
1678 fn find_verified_blob(&self, digest: &CacheDigest) -> Result<Option<PathBuf>> {
1679 let remembered = self.verified_blobs.lock().unwrap().get(digest).cloned();
1680 if let Some(remembered) = remembered {
1681 if remembered.is_unchanged() {
1692 return Ok(Some(remembered.path));
1693 }
1694 self.verified_blobs.lock().unwrap().remove(digest);
1695 }
1696 let path = self.cas.find(digest)?;
1697 if let Some(path) = &path {
1698 self.remember_verified_blob(digest, path);
1699 }
1700 Ok(path)
1701 }
1702
1703 fn remember_verified_blob(&self, digest: &CacheDigest, path: &Path) {
1704 let Some(verified) = VerifiedBlob::describe(path) else {
1705 return;
1709 };
1710 self.verified_blobs
1711 .lock()
1712 .unwrap()
1713 .insert(digest.clone(), verified);
1714 }
1715
1716 async fn find_action_result(&self, action: &CacheDigest) -> Result<AgentResponse> {
1717 if let Some(result) = self.actions.find(action)? {
1718 return Ok(AgentResponse::ActionResult {
1719 result: Some(result),
1720 });
1721 }
1722 if !self.remote_mode.reads() {
1723 return Ok(AgentResponse::ActionResult { result: None });
1724 }
1725 let Some(remote) = &self.remote else {
1726 return Ok(AgentResponse::ActionResult { result: None });
1727 };
1728 let lock = self.action_lock(action);
1729 let _guard = lock.lock().await;
1730 if let Some(result) = self.actions.find(action)? {
1731 return Ok(AgentResponse::ActionResult {
1732 result: Some(result),
1733 });
1734 }
1735 if let Some(result) = self
1736 .pending_remote_actions
1737 .lock()
1738 .unwrap()
1739 .get(action)
1740 .cloned()
1741 {
1742 return Ok(AgentResponse::ActionResult {
1743 result: Some(result),
1744 });
1745 }
1746 let _permit = self.remote_transfers.acquire().await?;
1747 match self.get_remote_action_result(remote, action).await {
1748 Ok(Some(result)) => {
1749 self.pending_remote_actions
1750 .lock()
1751 .unwrap()
1752 .insert(action.clone(), result.clone());
1753 Ok(AgentResponse::ActionResult {
1754 result: Some(result),
1755 })
1756 }
1757 Ok(None) => Ok(AgentResponse::ActionResult { result: None }),
1758 Err(error) => {
1759 self.note_remote_failure();
1760 warn!(
1761 "remote cache action lookup failed for {}: {error}",
1762 action.hash
1763 );
1764 Ok(AgentResponse::ActionResult { result: None })
1765 }
1766 }
1767 }
1768
1769 async fn store_action_result(
1770 &self,
1771 result: &RemoteActionResult,
1772 connection: &ConnectionUploads,
1773 ) -> Result<AgentResponse> {
1774 let path = self.actions.store(result)?;
1775 if let Some(uploads) = &self.uploads {
1776 uploads.queue_action_result(result, connection);
1777 }
1778 Ok(AgentResponse::ActionStored { path })
1779 }
1780
1781 async fn join_action_promise(
1782 &self,
1783 adapter: &str,
1784 invocation: &CacheDigest,
1785 ) -> Result<AgentResponse> {
1786 if !self.remote_mode.reads() || !self.remote_mode.writes() {
1787 return Ok(AgentResponse::ActionPromise {
1788 claim: None,
1789 prediction: None,
1790 });
1791 }
1792 let Some(remote) = &self.remote else {
1793 return Ok(AgentResponse::ActionPromise {
1794 claim: None,
1795 prediction: None,
1796 });
1797 };
1798 let deadline = Instant::now() + ACTION_PROMISE_WAIT;
1799 loop {
1800 let _permit = self.remote_transfers.acquire().await?;
1801 let state = remote.join_action_promise(invocation, adapter).await;
1802 drop(_permit);
1803 match state {
1804 Ok(Some(ActionPromiseState::Claimed { claim })) => {
1805 return Ok(AgentResponse::ActionPromise {
1806 claim: Some(claim),
1807 prediction: None,
1808 });
1809 }
1810 Ok(Some(ActionPromiseState::Complete { prediction })) => {
1811 return Ok(AgentResponse::ActionPromise {
1812 claim: None,
1813 prediction: Some(prediction),
1814 });
1815 }
1816 Ok(Some(ActionPromiseState::Pending { retry_after_ms }))
1817 if Instant::now() < deadline =>
1818 {
1819 tokio::time::sleep(Duration::from_millis(retry_after_ms.clamp(10, 5_000)))
1820 .await;
1821 }
1822 Ok(Some(ActionPromiseState::Pending { .. }) | None) => {
1823 return Ok(AgentResponse::ActionPromise {
1824 claim: None,
1825 prediction: None,
1826 });
1827 }
1828 Ok(Some(_)) => {
1829 return Ok(AgentResponse::ActionPromise {
1830 claim: None,
1831 prediction: None,
1832 });
1833 }
1834 Err(error) => {
1835 self.note_remote_failure();
1836 warn!(
1837 "remote cache action promise failed for {}: {error}",
1838 invocation.hash
1839 );
1840 return Ok(AgentResponse::ActionPromise {
1841 claim: None,
1842 prediction: None,
1843 });
1844 }
1845 }
1846 }
1847 }
1848
1849 async fn complete_action_promise(
1850 &self,
1851 claim: &str,
1852 prediction: &ActionPrediction,
1853 ) -> Result<AgentResponse> {
1854 prediction.validate()?;
1855 let Some(remote) = &self.remote else {
1856 return Ok(AgentResponse::ActionPromiseCompleted);
1857 };
1858 let Some(uploads) = &self.uploads else {
1859 return Ok(AgentResponse::ActionPromiseCompleted);
1860 };
1861 if uploads
1862 .wait_for_actions(std::slice::from_ref(&prediction.action))
1863 .await
1864 .contains(&prediction.action)
1865 {
1866 return Ok(AgentResponse::ActionPromiseCompleted);
1869 }
1870 let completion = ActionPromiseCompletion {
1871 claim: claim.to_string(),
1872 prediction: prediction.clone(),
1873 };
1874 let _permit = self.remote_transfers.acquire().await?;
1875 if let Err(error) = remote
1876 .complete_action_promise(&prediction.invocation, &completion)
1877 .await
1878 {
1879 self.note_remote_failure();
1880 warn!(
1881 "remote cache action promise completion failed for {}: {error}",
1882 prediction.invocation.hash
1883 );
1884 }
1885 Ok(AgentResponse::ActionPromiseCompleted)
1886 }
1887
1888 fn note_remote_failure(&self) {
1890 self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
1891 }
1892
1893 async fn get_remote_action_result(
1894 &self,
1895 remote: &RemoteCacheClient,
1896 action: &CacheDigest,
1897 ) -> Result<Option<RemoteActionResult>> {
1898 self.stats
1899 .remote_action_lookups
1900 .fetch_add(1, Ordering::Relaxed);
1901 let _timer = AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
1902 remote.get_action_result(action).await
1903 }
1904
1905 fn record_action_hit(
1906 &self,
1907 action: &CacheDigest,
1908 restore: RestoreStats,
1909 crate_name: Option<String>,
1910 diagnostic: Option<AgentEvent>,
1911 ) -> Result<AgentResponse> {
1912 validate_crate_name(crate_name.as_deref())?;
1913 if self.actions.find(action)?.is_none() {
1914 let pending = self.pending_remote_actions.lock().unwrap().remove(action);
1915 if let Some(result) = pending {
1916 self.actions.store(&result)?;
1917 } else {
1918 bail!("cannot record a hit for a missing action result");
1919 }
1920 }
1921 self.record_restore(restore);
1922 self.stats.hits.fetch_add(1, Ordering::Relaxed);
1923 self.emit_action(
1924 diagnostic,
1925 AgentEvent::ActionHit {
1926 crate_name,
1927 restore,
1928 },
1929 );
1930 Ok(AgentResponse::ActionHitRecorded)
1931 }
1932
1933 fn record_restore(&self, restore: RestoreStats) {
1934 self.record_materialization(restore);
1935 atomic_saturating_add(
1936 &self.stats.avoided_compiler_duration_ns,
1937 restore.avoided_compiler_duration_ns,
1938 );
1939 atomic_saturating_add(&self.stats.restored_output_files, restore.output_files);
1940 atomic_saturating_add(&self.stats.restored_output_bytes, restore.output_bytes);
1941 atomic_saturating_add(
1942 &self.stats.reflinked_output_files,
1943 restore.reflinked_output_files,
1944 );
1945 atomic_saturating_add(
1946 &self.stats.reflinked_output_bytes,
1947 restore.reflinked_output_bytes,
1948 );
1949 atomic_saturating_add(&self.stats.copied_output_files, restore.copied_output_files);
1950 atomic_saturating_add(&self.stats.copied_output_bytes, restore.copied_output_bytes);
1951 atomic_saturating_add(&self.stats.reused_output_files, restore.reused_output_files);
1952 atomic_saturating_add(&self.stats.reused_output_bytes, restore.reused_output_bytes);
1953 }
1954
1955 fn record_compiler_invocation(
1956 &self,
1957 outcome: &str,
1958 crate_name: Option<&str>,
1959 duration_ns: u64,
1960 diagnostic: Option<AgentEvent>,
1961 ) -> Result<AgentResponse> {
1962 if !matches!(
1963 outcome,
1964 "miss" | "unconsulted" | "bypass" | "verification" | "incremental"
1965 ) {
1966 bail!("invalid compiler invocation outcome");
1967 }
1968 validate_crate_name(crate_name)?;
1969 let mut compiler = self.stats.compiler.lock().unwrap();
1970 let stats = compiler.entry(outcome.to_string()).or_default();
1971 stats.invocations = stats.invocations.saturating_add(1);
1972 stats.duration_ns = stats.duration_ns.saturating_add(duration_ns);
1973 drop(compiler);
1974 if outcome != "verification"
1975 && let Some(crate_name) = crate_name.filter(|name| !name.is_empty())
1976 {
1977 let mut slow = self.stats.slow_compilations.lock().unwrap();
1978 let duration = slow.entry(crate_name.to_string()).or_default();
1979 *duration = duration.saturating_add(duration_ns);
1980 }
1981 self.emit_action(
1982 diagnostic,
1983 AgentEvent::CompilerInvocation {
1984 outcome: outcome.to_string(),
1985 crate_name: crate_name.map(str::to_string),
1986 duration_ns,
1987 },
1988 );
1989 Ok(AgentResponse::CompilerInvocationRecorded)
1990 }
1991
1992 fn record_materialization(&self, restore: RestoreStats) {
1993 atomic_saturating_add(&self.stats.materialization_duration_ns, restore.duration_ns);
1994 }
1995
1996 fn record_diagnostic(&self, level: DiagnosticLevel, message: String) -> Result<AgentResponse> {
2007 if message.is_empty()
2008 || message.len() > MAX_DIAGNOSTIC_BYTES
2009 || message.contains(['\n', '\r', '\0'])
2010 {
2011 bail!("invalid shim {}", level.label());
2012 }
2013 let mut diagnostics = self.diagnostics.lock().unwrap();
2014 let key = (level, message.clone());
2015 if !diagnostics.contains(&key) {
2016 if diagnostics.len() >= MAX_DIAGNOSTICS {
2020 bail!("shim diagnostic limit reached");
2021 }
2022 eprintln!("mbx[{}]: {message}", level.label());
2023 self.emit(|| match level {
2024 DiagnosticLevel::Warning => AgentEvent::Warning {
2025 message: message.clone(),
2026 },
2027 DiagnosticLevel::Error => AgentEvent::Error {
2028 message: message.clone(),
2029 },
2030 });
2031 diagnostics.insert(key);
2032 }
2033 Ok(match level {
2034 DiagnosticLevel::Warning => AgentResponse::WarningRecorded,
2035 DiagnosticLevel::Error => AgentResponse::ErrorRecorded,
2036 })
2037 }
2038
2039 fn find_file_digests(
2046 &self,
2047 scope: FileDigestScope,
2048 files: Vec<FileIdentity>,
2049 ) -> Result<AgentResponse> {
2050 if files.len() > MAX_FILE_DIGEST_BATCH {
2051 bail!("too many file-digest lookups in one request");
2052 }
2053 let ledger = self.file_digests.lock().unwrap();
2054 let digests = files
2055 .into_iter()
2056 .map(|file| {
2057 let recorded = ledger.get(&(scope, file.path.clone()))?;
2058 (recorded.file == file).then(|| recorded.digest.clone())
2059 })
2060 .collect();
2061 Ok(AgentResponse::FileDigests { digests })
2062 }
2063
2064 async fn resolve_file_digests(
2067 &self,
2068 scope: FileDigestScope,
2069 files: Vec<FileIdentity>,
2070 ) -> Result<AgentResponse> {
2071 if files.len() > MAX_FILE_DIGEST_BATCH {
2072 bail!("too many file-digest resolutions in one request");
2073 }
2074 for file in &files {
2075 if !file.path.is_absolute() {
2076 bail!("file-digest resolutions need absolute paths");
2077 }
2078 }
2079 let resolutions = stream::iter(
2080 files
2081 .into_iter()
2082 .map(|file| async move { self.resolve_file_digest(scope, file).await }),
2083 )
2084 .buffered(MAX_CONCURRENT_FILE_DIGESTS)
2085 .collect()
2086 .await;
2087 Ok(AgentResponse::FileDigestsResolved { resolutions })
2088 }
2089
2090 async fn resolve_file_digest(
2091 &self,
2092 scope: FileDigestScope,
2093 file: FileIdentity,
2094 ) -> FileDigestResolution {
2095 let lock = {
2096 let key = (scope, file.clone());
2097 let mut locks = self.file_digest_locks.lock().unwrap();
2098 locks.retain(|_, lock| lock.strong_count() > 0);
2099 if let Some(lock) = locks.get(&key).and_then(Weak::upgrade) {
2100 lock
2101 } else {
2102 let lock = Arc::new(FileDigestFlight {
2103 lock: tokio::sync::Mutex::new(()),
2104 resolution: Mutex::new(None),
2105 });
2106 locks.insert(key, Arc::downgrade(&lock));
2107 lock
2108 }
2109 };
2110 let _flight = lock.lock.lock().await;
2111 if let Some(resolution) = lock.resolution.lock().unwrap().clone() {
2112 return resolution;
2113 }
2114 if let Some(digest) = self
2115 .file_digests
2116 .lock()
2117 .unwrap()
2118 .get(&(scope, file.path.clone()))
2119 .filter(|recorded| recorded.file == file)
2120 .map(|recorded| recorded.digest.clone())
2121 {
2122 return FileDigestResolution::Digest(digest);
2123 }
2124 let Ok(_permit) = self.file_digest_permits.acquire().await else {
2125 return FileDigestResolution::Unresolved;
2126 };
2127 #[cfg(test)]
2128 self.file_digest_reads.fetch_add(1, Ordering::Relaxed);
2129 let path = file.path.clone();
2130 let resolved = tokio::task::spawn_blocking(move || {
2131 let resolution = digest_file(scope, &path)?;
2132 let current = std::fs::metadata(&path)
2133 .and_then(|metadata| FileIdentity::for_digest_cache(&path, &metadata));
2134 Ok::<_, std::io::Error>((resolution, current?))
2135 })
2136 .await;
2137 let Ok(Ok((resolution, current))) = resolved else {
2138 return FileDigestResolution::Unresolved;
2139 };
2140 let resolution = match resolution {
2141 FileDigestResolution::Digest(digest)
2142 if current.as_ref() == Some(&file) && digest.size == file.len =>
2143 {
2144 let _ = self.record_file_digests(
2145 scope,
2146 vec![RecordedFileDigest {
2147 file: file.clone(),
2148 digest: digest.clone(),
2149 }],
2150 );
2151 FileDigestResolution::Digest(digest)
2152 }
2153 FileDigestResolution::EmbeddedTimestampMacro if current.as_ref() == Some(&file) => {
2154 FileDigestResolution::EmbeddedTimestampMacro
2155 }
2156 FileDigestResolution::Digest(_)
2157 | FileDigestResolution::EmbeddedTimestampMacro
2158 | FileDigestResolution::Unresolved => FileDigestResolution::Unresolved,
2159 };
2160 *lock.resolution.lock().unwrap() = Some(resolution.clone());
2161 resolution
2162 }
2163
2164 fn record_file_digests(
2170 &self,
2171 scope: FileDigestScope,
2172 entries: Vec<RecordedFileDigest>,
2173 ) -> Result<AgentResponse> {
2174 if entries.len() > MAX_FILE_DIGEST_BATCH {
2175 bail!("too many file-digest records in one request");
2176 }
2177 for entry in &entries {
2178 validate_file_digest_record(entry)?;
2179 }
2180 let mut ledger = self.file_digests.lock().unwrap();
2181 for entry in entries {
2182 if ledger.len() >= MAX_FILE_DIGEST_ENTRIES
2183 && !ledger.contains_key(&(scope, entry.file.path.clone()))
2184 {
2185 break;
2186 }
2187 ledger.insert((scope, entry.file.path.clone()), entry);
2188 }
2189 Ok(AgentResponse::FileDigestsRecorded)
2190 }
2191
2192 pub fn seed_file_digests(&self, entries: Vec<(FileDigestScope, RecordedFileDigest)>) -> usize {
2201 self.seed_file_digests_with(|| entries)
2202 }
2203
2204 pub fn seed_file_digests_with(
2210 &self,
2211 load: impl FnOnce() -> Vec<(FileDigestScope, RecordedFileDigest)>,
2212 ) -> usize {
2213 let mut ledger = self.file_digests.lock().unwrap();
2214 let entries = load();
2215 let mut seeded = 0;
2216 for (scope, entry) in entries {
2217 if ledger.len() >= MAX_FILE_DIGEST_ENTRIES {
2218 break;
2219 }
2220 if validate_file_digest_record(&entry).is_err() {
2221 continue;
2222 }
2223 if let std::collections::btree_map::Entry::Vacant(slot) =
2224 ledger.entry((scope, entry.file.path.clone()))
2225 {
2226 slot.insert(entry);
2227 seeded += 1;
2228 }
2229 }
2230 seeded
2231 }
2232
2233 pub fn file_digests(&self) -> Vec<(FileDigestScope, RecordedFileDigest)> {
2236 self.file_digests
2237 .lock()
2238 .unwrap()
2239 .iter()
2240 .map(|((scope, _), entry)| (*scope, entry.clone()))
2241 .collect()
2242 }
2243
2244 fn find_action_prediction(
2245 &self,
2246 task: &str,
2247 invocation: &CacheDigest,
2248 ) -> Result<AgentResponse> {
2249 validate_task_identity(task)?;
2250 invocation.validate()?;
2251 let (prediction, prefetch) = {
2252 let mut tasks = self.task_actions.lock().unwrap();
2253 let Some(state) = tasks.get_mut(task) else {
2254 return Ok(AgentResponse::ActionPrediction { prediction: None });
2255 };
2256 activate_prediction_adapter(state, invocation)
2257 };
2258 if let Some(prefetch) = prefetch {
2259 self.spawn_prefetch_predictions(prefetch);
2260 }
2261 Ok(AgentResponse::ActionPrediction { prediction })
2262 }
2263
2264 fn record_action_prediction(
2265 &self,
2266 task: &str,
2267 prediction: ActionPrediction,
2268 ) -> Result<AgentResponse> {
2269 validate_task_identity(task)?;
2270 prediction.validate()?;
2271 let mut tasks = self.task_actions.lock().unwrap();
2272 let state = tasks.entry(task.to_string()).or_default();
2273 if !state.predictions.contains_key(&prediction.invocation)
2274 && state.predictions.len() >= MAX_TASK_ACTION_PREDICTIONS
2275 {
2276 bail!("task action manifest contains too many predictions");
2277 }
2278 let current = state.predictions.get(&prediction.invocation);
2279 if current != Some(&prediction) {
2280 match state.changed_predictions.get(&prediction.invocation) {
2281 Some(baseline) if baseline.as_ref() == Some(&prediction) => {
2283 state.changed_predictions.remove(&prediction.invocation);
2284 }
2285 Some(_) => {}
2287 None => {
2288 state
2289 .changed_predictions
2290 .insert(prediction.invocation.clone(), current.cloned());
2291 }
2292 }
2293 }
2294 state
2295 .predictions
2296 .insert(prediction.invocation.clone(), prediction.clone());
2297 state
2298 .pending_predictions
2299 .insert(prediction.invocation.clone(), prediction);
2300 Ok(AgentResponse::ActionPredictionRecorded)
2301 }
2302
2303 fn executable_identity_key(
2304 &self,
2305 executable: PathBuf,
2306 environment: BTreeMap<String, Option<String>>,
2307 ) -> Result<ExecutableIdentityKey> {
2308 if !environment.keys().all(|name| {
2314 matches!(
2315 name.as_str(),
2316 "COMPILER_PATH"
2317 | "GCC_EXEC_PREFIX"
2318 | "LIBRARY_PATH"
2319 | "MBX_FUSE_LD"
2320 | "PATH"
2321 | "RUSTUP_HOME"
2322 | "RUSTUP_TOOLCHAIN"
2323 | "SDKROOT"
2324 | "MACOSX_DEPLOYMENT_TARGET"
2325 | "LIB"
2326 | "UCRTVersion"
2327 | "UniversalCRTSdkDir"
2328 | "VCToolsInstallDir"
2329 | "VCToolsVersion"
2330 | "WindowsSdkDir"
2331 | "WindowsSDKVersion"
2332 )
2333 }) {
2334 bail!("executable identity contains an unsupported environment variable");
2335 }
2336 Ok(ExecutableIdentityKey {
2337 executable,
2338 environment,
2339 })
2340 }
2341
2342 fn find_executable_identity(
2343 &self,
2344 executable: PathBuf,
2345 environment: BTreeMap<String, Option<String>>,
2346 ) -> Result<AgentResponse> {
2347 let key = self.executable_identity_key(executable, environment)?;
2348 let remembered = self
2349 .executable_identities
2350 .lock()
2351 .unwrap()
2352 .get(&key)
2353 .cloned();
2354 let stdout = match remembered {
2355 Some(stdout) => Some(stdout),
2356 None => {
2357 let persisted = self.load_persisted_identity(&key);
2358 if let Some(stdout) = &persisted {
2359 let _ = self.remember_executable_identity(&key, stdout.clone());
2362 }
2363 persisted
2364 }
2365 };
2366 Ok(AgentResponse::ExecutableIdentity { stdout })
2367 }
2368
2369 fn store_executable_identity(
2370 &self,
2371 executable: PathBuf,
2372 environment: BTreeMap<String, Option<String>>,
2373 stdout: Vec<u8>,
2374 pins: Vec<PinnedFile>,
2375 ) -> Result<AgentResponse> {
2376 if stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE {
2377 bail!("executable identity exceeds {MAX_EXECUTABLE_IDENTITY_SIZE} bytes");
2378 }
2379 let key = self.executable_identity_key(executable, environment)?;
2380 self.remember_executable_identity(&key, stdout.clone())?;
2381 if pins.is_empty() {
2386 } else if !pins.iter().all(PinnedFile::holds) {
2387 debug!("executable identity was not persisted: a pinned file changed during the probe");
2388 } else if let Err(error) = self.persist_executable_identity(&key, pins, &stdout) {
2389 debug!("executable identity was not persisted: {error:#}");
2390 }
2391 Ok(AgentResponse::ExecutableIdentity {
2392 stdout: Some(stdout),
2393 })
2394 }
2395
2396 fn remember_executable_identity(
2397 &self,
2398 key: &ExecutableIdentityKey,
2399 stdout: Vec<u8>,
2400 ) -> Result<()> {
2401 let mut identities = self.executable_identities.lock().unwrap();
2402 let is_new = !identities.contains_key(key);
2403 let previous_size = identities.get(key).map_or(0, Vec::len);
2404 if is_new && identities.len() >= MAX_EXECUTABLE_IDENTITIES {
2405 bail!("executable identity cache contains too many entries");
2406 }
2407 let retained_bytes = identities.values().map(Vec::len).sum::<usize>();
2408 if retained_bytes - previous_size + stdout.len() > MAX_EXECUTABLE_IDENTITY_BYTES {
2409 bail!("executable identity cache contains too many bytes");
2410 }
2411 identities.insert(key.clone(), stdout);
2412 Ok(())
2413 }
2414
2415 fn persisted_identity_path(&self, key: &ExecutableIdentityKey) -> Result<PathBuf> {
2416 let selector = serde_json::to_vec(&(&key.executable, &key.environment))?;
2417 let digest = CacheDigest::blake3(&selector);
2418 Ok(self.identity_dir.join(format!("{}.json", digest.hash)))
2419 }
2420
2421 fn load_persisted_identity(&self, key: &ExecutableIdentityKey) -> Option<Vec<u8>> {
2424 let path = self.persisted_identity_path(key).ok()?;
2425 let bytes = fs::read(path).ok()?;
2426 let persisted: PersistedExecutableIdentity = serde_json::from_slice(&bytes).ok()?;
2427 if persisted.version != PERSISTED_EXECUTABLE_IDENTITY_VERSION
2428 || persisted.executable != key.executable
2429 || persisted.environment != key.environment
2430 || persisted.pins.is_empty()
2431 || persisted.stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE
2432 {
2433 return None;
2434 }
2435 persisted
2436 .pins
2437 .iter()
2438 .all(PinnedFile::holds)
2439 .then(|| persisted.stdout.into_bytes())
2440 }
2441
2442 fn persist_executable_identity(
2443 &self,
2444 key: &ExecutableIdentityKey,
2445 pins: Vec<PinnedFile>,
2446 stdout: &[u8],
2447 ) -> Result<()> {
2448 let Ok(stdout) = std::str::from_utf8(stdout) else {
2450 return Ok(());
2451 };
2452 let persisted = PersistedExecutableIdentity {
2453 version: PERSISTED_EXECUTABLE_IDENTITY_VERSION,
2454 executable: key.executable.clone(),
2455 environment: key.environment.clone(),
2456 pins,
2457 stdout: stdout.to_owned(),
2458 };
2459 let path = self.persisted_identity_path(key)?;
2460 fs::create_dir_all(self.identity_dir.as_path())?;
2461 let mut temporary = tempfile::NamedTempFile::new_in(self.identity_dir.as_path())?;
2463 std::io::Write::write_all(temporary.as_file_mut(), &serde_json::to_vec(&persisted)?)?;
2464 temporary.persist(path).map_err(|error| error.error)?;
2465 Ok(())
2466 }
2467
2468 pub async fn handle_connection<S>(&self, stream: S) -> Result<()>
2470 where
2471 S: AsyncRead + AsyncWrite + Unpin,
2472 {
2473 let (reader, mut writer) = tokio::io::split(stream);
2474 let mut reader = BufReader::new(reader);
2475 let hello = read_request(&mut reader)
2476 .await?
2477 .ok_or_else(|| eyre::eyre!("connection closed before the agent handshake"))?;
2478 let request: AgentRequest = serde_json::from_str(&hello)?;
2479 match request {
2480 AgentRequest::Hello {
2481 protocol,
2482 client_version,
2483 } if protocol == AGENT_PROTOCOL_VERSION && client_version == self.version.as_ref() => {}
2484 AgentRequest::Hello { protocol, .. } if protocol != AGENT_PROTOCOL_VERSION => {
2485 send_response(
2486 &mut writer,
2487 &AgentResponse::Error {
2488 message: format!(
2489 "unsupported agent protocol {protocol}; expected {AGENT_PROTOCOL_VERSION}"
2490 ),
2491 },
2492 )
2493 .await?;
2494 return Ok(());
2495 }
2496 AgentRequest::Hello { client_version, .. } => {
2497 send_response(
2498 &mut writer,
2499 &AgentResponse::Error {
2500 message: format!(
2501 "cache client {client_version} does not match agent {}",
2502 self.version
2503 ),
2504 },
2505 )
2506 .await?;
2507 return Ok(());
2508 }
2509 _ => bail!("the first agent request must be hello"),
2510 }
2511 send_response(
2512 &mut writer,
2513 &AgentResponse::Hello {
2514 protocol: AGENT_PROTOCOL_VERSION,
2515 agent_version: self.version.to_string(),
2516 },
2517 )
2518 .await?;
2519
2520 let mut connection = ConnectionUploads::default();
2524 while let Some(line) = read_request(&mut reader).await? {
2525 let response = match serde_json::from_str(&line) {
2526 Ok(request) => self.respond_on(request, &mut connection).await,
2527 Err(error) => AgentResponse::Error {
2528 message: format!("invalid agent request: {error}"),
2529 },
2530 };
2531 send_response(&mut writer, &response).await?;
2532 }
2533 Ok(())
2534 }
2535}
2536
2537async fn read_request<R>(reader: &mut R) -> Result<Option<String>>
2543where
2544 R: AsyncBufRead + Unpin,
2545{
2546 let mut line = Vec::new();
2547 loop {
2548 let available = reader.fill_buf().await?;
2549 if available.is_empty() {
2550 break;
2551 }
2552 let (consumed, complete) = match available.iter().position(|byte| *byte == b'\n') {
2553 Some(index) => (index, true),
2554 None => (available.len(), false),
2555 };
2556 if line.len() + consumed > MAX_REQUEST_BYTES {
2557 bail!("agent request exceeded {MAX_REQUEST_BYTES} bytes");
2558 }
2559 line.extend_from_slice(&available[..consumed]);
2560 reader.consume(consumed + usize::from(complete));
2562 if complete {
2563 return Ok(Some(String::from_utf8(line)?));
2564 }
2565 }
2566 if line.is_empty() {
2567 Ok(None)
2568 } else {
2569 Ok(Some(String::from_utf8(line)?))
2570 }
2571}
2572
2573async fn send_response(
2574 writer: &mut (impl AsyncWrite + Unpin),
2575 response: &AgentResponse,
2576) -> Result<()> {
2577 let mut encoded = serde_json::to_vec(response)?;
2578 encoded.push(b'\n');
2579 writer.write_all(&encoded).await?;
2580 writer.flush().await?;
2581 Ok(())
2582}
2583
2584#[cfg(test)]
2585#[path = "agent_tests.rs"]
2586mod tests;