1use crate::uploads::{ConnectionUploads, UploadQueue, UploadSink};
2use crate::{
3 ActionPrediction, BlobPackLimits, CacheDigest, CacheDirectory, LocalActionCache, LocalCas,
4 MAX_STAGED_BLOB_PACK_BYTES, MAX_STAGED_BLOB_PACK_ITEMS, 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::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};
18use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
19
20const MAX_EXECUTABLE_IDENTITIES: usize = 64;
21const MAX_EXECUTABLE_IDENTITY_SIZE: usize = 64 * 1024;
22const MAX_EXECUTABLE_IDENTITY_BYTES: usize = 256 * 1024;
23const TASK_ACTION_MANIFEST_VERSION: u8 = 1;
24const MAX_TASK_ACTION_PREDICTIONS: usize = 16 * 1024;
25const MAX_WARNING_BYTES: usize = 4 * 1024;
27const MAX_WARNINGS: usize = 128;
33const MAX_REMOTE_TRANSFERS: usize = 64;
34const MAX_PREFETCH_TRANSFERS: usize = 48;
35const MAX_PREFETCH_ACTION_BATCH: usize = 256;
36const MAX_PREFETCH_BATCH_LOOKUPS: usize = 4;
39const PREFETCH_ACTION_BATCH_DELAY: Duration = Duration::from_millis(5);
40const MAX_PREFETCH_DIRECTORY_OBJECTS: usize = 100_000;
41const MAX_PREFETCH_OBJECTS_PER_WAVE: usize = 100_000;
42const DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES: u64 = 5 * 1024 * 1024 * 1024;
43
44pub struct AgentRemoteCache {
46 pub client: RemoteCacheClient,
48 pub mode: RemoteCacheMode,
50 pub staging_dir: PathBuf,
52}
53
54pub const AGENT_PROTOCOL_VERSION: u8 = 4;
56const MAX_REQUEST_BYTES: usize = 16 * 1024 * 1024;
61
62#[derive(Debug, Clone, Serialize, Deserialize)]
64#[serde(tag = "type", rename_all = "snake_case")]
65pub enum AgentRequest {
66 Hello {
68 protocol: u8,
70 client_version: String,
72 },
73 BeginTask {
75 task: String,
77 },
78 CommitTask {
80 run: String,
82 },
83 FindBlob {
85 digest: CacheDigest,
87 },
88 FindBlobs {
90 digests: Vec<CacheDigest>,
92 },
93 StoreBlob {
95 digest: CacheDigest,
97 source: PathBuf,
99 },
100 FindActionResult {
102 action: CacheDigest,
104 },
105 RecordActionHit {
107 action: CacheDigest,
109 restore: RestoreStats,
111 #[serde(default)]
113 crate_name: Option<String>,
114 },
115 RecordBypass {
117 kind: String,
119 },
120 RecordUnconsulted,
123 RecordCompilerInvocation {
125 outcome: String,
129 crate_name: Option<String>,
131 duration_ns: u64,
133 },
134 RecordActionVerification {
136 matched: bool,
138 restore: RestoreStats,
140 },
141 StoreActionResult {
143 result: RemoteActionResult,
145 },
146 FindActionPrediction {
148 task: String,
150 invocation: CacheDigest,
152 },
153 RecordActionPrediction {
155 task: String,
157 prediction: ActionPrediction,
159 },
160 FindExecutableIdentity {
162 executable: PathBuf,
164 environment: BTreeMap<String, Option<String>>,
166 },
167 StoreExecutableIdentity {
169 executable: PathBuf,
171 environment: BTreeMap<String, Option<String>>,
173 stdout: Vec<u8>,
175 },
176 RecordWarning {
189 message: String,
191 },
192}
193
194#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
196#[serde(deny_unknown_fields)]
197pub struct RestoreStats {
198 pub duration_ns: u64,
200 pub avoided_compiler_duration_ns: u64,
203 pub output_files: u64,
205 pub output_bytes: u64,
207 pub reflinked_output_files: u64,
209 pub reflinked_output_bytes: u64,
211 pub copied_output_files: u64,
213 pub copied_output_bytes: u64,
215}
216
217#[derive(Debug, Clone)]
224#[non_exhaustive]
225pub enum AgentEvent {
226 ActionHit {
228 crate_name: Option<String>,
230 restore: RestoreStats,
232 },
233 Bypass {
235 kind: String,
237 },
238 Unconsulted,
240 CompilerInvocation {
242 outcome: String,
244 crate_name: Option<String>,
246 duration_ns: u64,
248 },
249 Verification {
251 matched: bool,
253 restore: RestoreStats,
255 },
256 Warning {
258 message: String,
260 },
261}
262
263pub trait AgentEventObserver: Send + Sync {
265 fn event(&self, event: AgentEvent);
267}
268
269#[derive(Debug, Clone, Serialize, Deserialize)]
271#[serde(tag = "type", rename_all = "snake_case")]
272pub enum AgentResponse {
273 Hello {
275 protocol: u8,
277 agent_version: String,
279 },
280 TaskBegun {
282 run: String,
284 },
285 TaskCommitted,
287 Blob {
289 path: Option<PathBuf>,
291 },
292 Blobs {
294 paths: Vec<Option<PathBuf>>,
296 },
297 Stored {
299 path: PathBuf,
301 },
302 ActionResult {
304 result: Option<RemoteActionResult>,
306 },
307 ActionHitRecorded,
309 ActionVerificationRecorded,
311 BypassRecorded,
313 UnconsultedRecorded,
315 CompilerInvocationRecorded,
317 ActionStored {
319 path: PathBuf,
321 },
322 ActionPrediction {
324 prediction: Option<ActionPrediction>,
326 },
327 ActionPredictionRecorded,
329 ExecutableIdentity {
331 stdout: Option<Vec<u8>>,
333 },
334 Error {
336 message: String,
338 },
339 WarningRecorded,
345}
346
347#[derive(Debug, Clone, Default, PartialEq, Eq)]
354#[non_exhaustive]
355pub struct AgentStats {
356 pub session_duration_ns: u64,
358 pub lookups: u64,
360 pub unconsulted: u64,
366 pub hits: u64,
368 pub stores: u64,
370 pub stored_bytes: u64,
372 pub verifications: u64,
374 pub divergences: u64,
376 pub downloaded_bytes: u64,
378 pub uploaded_bytes: u64,
380 pub background_uploads: u64,
386 pub background_upload_failures: u64,
389 pub remote_blob_pack_uploads: u64,
391 pub remote_blob_pack_upload_blobs: u64,
393 pub upload_drain_duration_ns: u64,
395 pub prefetched_actions: u64,
397 pub predictions_loaded: u64,
406 pub bypasses: BTreeMap<String, u64>,
408 pub avoided_compiler_duration_ns: u64,
410 pub compiler: BTreeMap<String, CompilerStats>,
412 pub slow_compilations: BTreeMap<String, u64>,
414 pub remote_failures: u64,
422 pub remote_manifest_lookups: u64,
424 pub remote_manifest_lookup_duration_ns: u64,
426 pub remote_action_lookups: u64,
428 pub remote_action_lookup_duration_ns: u64,
430 pub remote_blob_requests: u64,
432 pub remote_blob_pack_requests: u64,
434 pub remote_blob_pack_blobs: u64,
436 pub remote_blob_transfer_duration_ns: u64,
438 pub local_cas_write_duration_ns: u64,
440 pub prefetch_runs: u64,
442 pub prefetch_duration_ns: u64,
444 pub materialization_duration_ns: u64,
446 pub restored_output_files: u64,
448 pub restored_output_bytes: u64,
450 pub reflinked_output_files: u64,
452 pub reflinked_output_bytes: u64,
454 pub copied_output_files: u64,
456 pub copied_output_bytes: u64,
458}
459
460#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
466#[non_exhaustive]
467pub struct CompilerStats {
468 pub invocations: u64,
470 pub duration_ns: u64,
472}
473
474impl CompilerStats {
475 pub fn new(invocations: u64, duration_ns: u64) -> Self {
477 Self {
478 invocations,
479 duration_ns,
480 }
481 }
482}
483
484#[derive(Default)]
485struct AtomicAgentStats {
486 lookups: AtomicU64,
487 unconsulted: AtomicU64,
488 hits: AtomicU64,
489 stores: AtomicU64,
490 stored_bytes: AtomicU64,
491 verifications: AtomicU64,
492 divergences: AtomicU64,
493 downloaded_bytes: AtomicU64,
494 uploaded_bytes: AtomicU64,
495 background_uploads: AtomicU64,
496 background_upload_failures: AtomicU64,
497 remote_blob_pack_uploads: AtomicU64,
498 remote_blob_pack_upload_blobs: AtomicU64,
499 upload_drain_duration_ns: AtomicU64,
500 prefetched_actions: AtomicU64,
501 predictions_loaded: AtomicU64,
502 remote_failures: AtomicU64,
503 remote_manifest_lookups: AtomicU64,
504 remote_manifest_lookup_duration_ns: AtomicU64,
505 remote_action_lookups: AtomicU64,
506 remote_action_lookup_duration_ns: AtomicU64,
507 remote_blob_requests: AtomicU64,
508 remote_blob_pack_requests: AtomicU64,
509 remote_blob_pack_blobs: AtomicU64,
510 remote_blob_transfer_duration_ns: AtomicU64,
511 local_cas_write_duration_ns: AtomicU64,
512 prefetch_runs: AtomicU64,
513 prefetch_duration_ns: AtomicU64,
514 materialization_duration_ns: AtomicU64,
515 bypasses: Mutex<BTreeMap<String, u64>>,
516 avoided_compiler_duration_ns: AtomicU64,
517 compiler: Mutex<BTreeMap<String, CompilerStats>>,
518 slow_compilations: Mutex<BTreeMap<String, u64>>,
519 restored_output_files: AtomicU64,
520 restored_output_bytes: AtomicU64,
521 reflinked_output_files: AtomicU64,
522 reflinked_output_bytes: AtomicU64,
523 copied_output_files: AtomicU64,
524 copied_output_bytes: AtomicU64,
525}
526
527struct AtomicDurationTimer<'a> {
528 started: Instant,
529 target: &'a AtomicU64,
530}
531
532impl<'a> AtomicDurationTimer<'a> {
533 fn start(target: &'a AtomicU64) -> Self {
534 Self {
535 started: Instant::now(),
536 target,
537 }
538 }
539}
540
541impl Drop for AtomicDurationTimer<'_> {
542 fn drop(&mut self) {
543 atomic_saturating_add(self.target, duration_ns(self.started));
544 }
545}
546
547fn duration_ns(started: Instant) -> u64 {
548 started.elapsed().as_nanos().try_into().unwrap_or(u64::MAX)
549}
550
551fn validate_crate_name(crate_name: Option<&str>) -> Result<()> {
552 if let Some(crate_name) = crate_name
553 && (crate_name.len() > 256 || crate_name.contains(['\0', '\n', '\r']))
554 {
555 bail!("invalid compiler crate name");
556 }
557 Ok(())
558}
559
560fn atomic_saturating_add(target: &AtomicU64, value: u64) {
561 let _ = target.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
562 Some(current.saturating_add(value))
563 });
564}
565
566fn queue_prefetch_digest(
567 verified: &BTreeMap<CacheDigest, PathBuf>,
568 pending: &mut BTreeMap<CacheDigest, ()>,
569 digest: CacheDigest,
570) {
571 if verified.contains_key(&digest) || pending.contains_key(&digest) {
572 return;
573 }
574 pending.insert(digest, ());
575}
576
577fn queue_prefetch_directory(
578 seen: &BTreeMap<CacheDigest, ()>,
579 pending: &mut BTreeMap<CacheDigest, ()>,
580 digest: CacheDigest,
581 limit: usize,
582) -> bool {
583 if seen.contains_key(&digest) || pending.contains_key(&digest) {
584 return true;
585 }
586 if seen.len().saturating_add(pending.len()) >= limit {
587 return false;
588 }
589 pending.insert(digest, ());
590 true
591}
592
593#[derive(Clone)]
599pub struct CacheAgent {
600 cas: LocalCas,
601 actions: LocalActionCache,
602 verified_blobs: Arc<Mutex<BTreeMap<CacheDigest, VerifiedBlob>>>,
603 version: Arc<str>,
604 write_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
605 action_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
606 stats: Arc<AtomicAgentStats>,
607 observer: Option<Arc<dyn AgentEventObserver>>,
608 executable_identities: Arc<Mutex<BTreeMap<ExecutableIdentityKey, Vec<u8>>>>,
609 manifest_dir: Arc<PathBuf>,
610 task_actions: Arc<Mutex<BTreeMap<String, TaskActionState>>>,
611 next_task_run: Arc<AtomicU64>,
612 manifest_write_lock: Arc<Mutex<()>>,
613 remote: Option<Arc<RemoteCacheClient>>,
614 remote_mode: RemoteCacheMode,
615 remote_staging_dir: Arc<PathBuf>,
616 remote_download_limit: u64,
617 remote_download_bytes: Arc<AtomicU64>,
618 pending_remote_actions: Arc<Mutex<BTreeMap<CacheDigest, RemoteActionResult>>>,
619 remote_transfers: Arc<tokio::sync::Semaphore>,
620 prefetch_transfers: Arc<tokio::sync::Semaphore>,
621 prefetch_tasks: Arc<Mutex<Vec<tokio::task::JoinHandle<()>>>>,
622 warnings: Arc<Mutex<BTreeSet<String>>>,
625 uploads: Option<UploadQueue>,
627}
628
629struct AgentUploadSink {
631 stats: Arc<AtomicAgentStats>,
632}
633
634impl UploadSink for AgentUploadSink {
635 fn record_blob_uploaded(&self, bytes: u64) {
636 self.stats
637 .background_uploads
638 .fetch_add(1, Ordering::Relaxed);
639 self.stats
640 .uploaded_bytes
641 .fetch_add(bytes, Ordering::Relaxed);
642 }
643
644 fn record_action_uploaded(&self) {
645 self.stats
646 .background_uploads
647 .fetch_add(1, Ordering::Relaxed);
648 }
649
650 fn record_blob_pack_uploaded(&self, blobs: u64) {
651 self.stats
652 .remote_blob_pack_uploads
653 .fetch_add(1, Ordering::Relaxed);
654 self.stats
655 .remote_blob_pack_upload_blobs
656 .fetch_add(blobs, Ordering::Relaxed);
657 }
658
659 fn record_upload_failure(&self) {
660 self.stats
661 .background_upload_failures
662 .fetch_add(1, Ordering::Relaxed);
663 self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
664 }
665}
666
667#[derive(Debug, Clone)]
673struct VerifiedBlob {
674 path: PathBuf,
675 len: u64,
676 modified: SystemTime,
677}
678
679impl VerifiedBlob {
680 fn describe(path: &Path) -> Option<Self> {
683 let metadata = std::fs::metadata(path).ok()?;
684 Some(Self {
685 path: path.to_path_buf(),
686 len: metadata.len(),
687 modified: metadata.modified().ok()?,
688 })
689 }
690
691 fn is_unchanged(&self) -> bool {
693 let Ok(metadata) = std::fs::metadata(&self.path) else {
694 return false;
695 };
696 metadata.len() == self.len && metadata.modified().is_ok_and(|now| now == self.modified)
697 }
698}
699
700#[derive(Debug, Clone, Default)]
701struct TaskActionState {
702 manifest: String,
703 baseline_loaded: bool,
704 predictions: BTreeMap<CacheDigest, ActionPrediction>,
705 pending_predictions: BTreeMap<CacheDigest, ActionPrediction>,
706 remote_etag: Option<String>,
707}
708
709struct PrefetchedAction {
710 adapter: String,
711 result: RemoteActionResult,
712}
713
714struct RemoteDownloadReservation {
715 counter: Arc<AtomicU64>,
716 reserved: u64,
717 committed: bool,
718}
719
720impl RemoteDownloadReservation {
721 fn bytes(&self) -> u64 {
722 self.reserved
723 }
724
725 fn commit(mut self, bytes: u64) {
726 debug_assert!(bytes <= self.reserved);
727 self.counter
728 .fetch_sub(self.reserved.saturating_sub(bytes), Ordering::AcqRel);
729 self.committed = true;
730 }
731}
732
733impl Drop for RemoteDownloadReservation {
734 fn drop(&mut self) {
735 if !self.committed {
736 self.counter.fetch_sub(self.reserved, Ordering::AcqRel);
737 }
738 }
739}
740
741#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
742struct ExecutableIdentityKey {
743 executable: PathBuf,
744 environment: BTreeMap<String, Option<String>>,
745}
746
747impl CacheAgent {
748 pub fn new(cache_dir: impl Into<PathBuf>, version: impl Into<Arc<str>>) -> Self {
750 Self::build(cache_dir.into(), version.into(), None, 0)
751 }
752
753 pub fn new_remote(
755 cache_dir: impl Into<PathBuf>,
756 version: impl Into<Arc<str>>,
757 remote: AgentRemoteCache,
758 ) -> Self {
759 Self::build(
760 cache_dir.into(),
761 version.into(),
762 Some(remote),
763 DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES,
764 )
765 }
766
767 pub fn new_remote_with_download_limit(
769 cache_dir: impl Into<PathBuf>,
770 version: impl Into<Arc<str>>,
771 remote: AgentRemoteCache,
772 max_remote_download_bytes: u64,
773 ) -> Self {
774 Self::build(
775 cache_dir.into(),
776 version.into(),
777 Some(remote),
778 max_remote_download_bytes,
779 )
780 }
781
782 fn build(
783 cache_dir: PathBuf,
784 version: Arc<str>,
785 remote: Option<AgentRemoteCache>,
786 remote_download_limit: u64,
787 ) -> Self {
788 let remote_mode = remote
789 .as_ref()
790 .map_or(RemoteCacheMode::ReadOnly, |remote| remote.mode);
791 let remote_staging_dir = remote.as_ref().map_or_else(
792 || cache_dir.join("remote"),
793 |remote| remote.staging_dir.clone(),
794 );
795 let remote = remote.map(|remote| Arc::new(remote.client));
796 let stats = Arc::new(AtomicAgentStats::default());
797 let remote_transfers = Arc::new(tokio::sync::Semaphore::new(MAX_REMOTE_TRANSFERS));
798 let uploads = remote
801 .clone()
802 .filter(|_| remote_mode.writes())
803 .map(|client| {
804 UploadQueue::new(
805 client,
806 Arc::new(AgentUploadSink {
807 stats: stats.clone(),
808 }),
809 remote_transfers.clone(),
810 )
811 });
812 Self {
813 cas: LocalCas::new(cache_dir.clone()),
814 actions: LocalActionCache::new(cache_dir.clone()),
815 verified_blobs: Arc::new(Mutex::new(BTreeMap::new())),
816 version,
817 write_locks: Arc::new(Mutex::new(BTreeMap::new())),
818 action_locks: Arc::new(Mutex::new(BTreeMap::new())),
819 stats,
820 observer: None,
821 executable_identities: Arc::new(Mutex::new(BTreeMap::new())),
822 manifest_dir: Arc::new(task_manifest_dir(&cache_dir)),
823 task_actions: Arc::new(Mutex::new(BTreeMap::new())),
824 next_task_run: Arc::new(AtomicU64::new(0)),
825 manifest_write_lock: Arc::new(Mutex::new(())),
826 remote,
827 remote_mode,
828 remote_staging_dir: Arc::new(remote_staging_dir),
829 remote_download_limit,
830 remote_download_bytes: Arc::new(AtomicU64::new(0)),
831 pending_remote_actions: Arc::new(Mutex::new(BTreeMap::new())),
832 remote_transfers,
833 prefetch_transfers: Arc::new(tokio::sync::Semaphore::new(MAX_PREFETCH_TRANSFERS)),
834 prefetch_tasks: Arc::new(Mutex::new(Vec::new())),
835 warnings: Arc::new(Mutex::new(BTreeSet::new())),
836 uploads,
837 }
838 }
839
840 #[must_use]
842 pub fn with_observer(mut self, observer: Arc<dyn AgentEventObserver>) -> Self {
843 self.observer = Some(observer);
844 self
845 }
846
847 fn emit(&self, event: impl FnOnce() -> AgentEvent) {
848 if let Some(observer) = &self.observer {
849 observer.event(event());
850 }
851 }
852
853 fn reserve_remote_download(&self, bytes: u64) -> Result<RemoteDownloadReservation> {
854 let mut current = self.remote_download_bytes.load(Ordering::Acquire);
855 loop {
856 let next = current
857 .checked_add(bytes)
858 .ok_or_else(|| eyre::eyre!("remote cache download budget overflowed"))?;
859 if next > self.remote_download_limit {
860 bail!(
861 "remote cache download budget exceeded: {} bytes requested with {} of {} bytes already used",
862 bytes,
863 current,
864 self.remote_download_limit
865 );
866 }
867 match self.remote_download_bytes.compare_exchange_weak(
868 current,
869 next,
870 Ordering::AcqRel,
871 Ordering::Acquire,
872 ) {
873 Ok(_) => {
874 return Ok(RemoteDownloadReservation {
875 counter: self.remote_download_bytes.clone(),
876 reserved: bytes,
877 committed: false,
878 });
879 }
880 Err(observed) => current = observed,
881 }
882 }
883 }
884
885 fn reserve_remote_download_up_to(&self, requested: u64) -> Result<RemoteDownloadReservation> {
886 let mut current = self.remote_download_bytes.load(Ordering::Acquire);
887 loop {
888 let reserved = requested.min(self.remote_download_limit.saturating_sub(current));
889 let next = current + reserved;
890 match self.remote_download_bytes.compare_exchange_weak(
891 current,
892 next,
893 Ordering::AcqRel,
894 Ordering::Acquire,
895 ) {
896 Ok(_) => {
897 return Ok(RemoteDownloadReservation {
898 counter: self.remote_download_bytes.clone(),
899 reserved,
900 committed: false,
901 });
902 }
903 Err(observed) => current = observed,
904 }
905 }
906 }
907
908 pub async fn begin_task(&self, task: &str) -> Result<String> {
910 self.begin_task_with_remote_errors(task, false).await
911 }
912
913 pub async fn prefetch_task(&self, task: &str) -> Result<String> {
915 let run = self.begin_task_with_remote_errors(task, true).await?;
916 self.wait_for_prefetches().await;
917 Ok(run)
918 }
919
920 async fn begin_task_with_remote_errors(&self, task: &str, strict: bool) -> Result<String> {
921 validate_task_identity(task)?;
922 let (remote_manifest, mut remote_etag) = if self.remote_mode.reads() {
923 match self.get_remote_task_manifest(task).await {
924 Ok(Some((manifest, etag))) => (Some(manifest), Some(etag)),
925 Ok(None) => (None, None),
926 Err(error) => {
927 if strict {
928 return Err(error).wrap_err_with(|| {
929 format!("remote task action manifest lookup failed for {task}")
930 });
931 }
932 self.note_remote_failure();
933 warn!("remote task action manifest lookup failed for {task}: {error}");
934 (None, None)
935 }
936 }
937 } else {
938 (None, None)
939 };
940 let manifest = {
941 let _write_guard = self.manifest_write_lock.lock().unwrap();
942 let _file_guard = self.lock_task_manifest(task)?;
943 let local_manifest = self.load_task_manifest(task)?;
944 let manifest = match (remote_manifest, local_manifest) {
945 (Some(remote), Some(local)) => {
946 let (manifest, merged) = merge_remote_task_manifest(task, remote, local);
947 if !merged {
948 remote_etag = None;
949 }
950 Some(manifest)
951 }
952 (Some(remote), None) => Some(remote),
953 (None, local) => local,
954 };
955 if let Some(manifest) = &manifest {
956 self.persist_task_manifest(manifest)?;
957 }
958 manifest
959 };
960 let state = if let Some(manifest) = manifest {
961 TaskActionState {
962 manifest: task.to_string(),
963 baseline_loaded: true,
964 predictions: manifest
965 .predictions
966 .into_iter()
967 .map(|prediction| (prediction.invocation.clone(), prediction))
968 .collect(),
969 pending_predictions: BTreeMap::new(),
970 remote_etag,
971 }
972 } else {
973 TaskActionState {
974 manifest: task.to_string(),
975 baseline_loaded: true,
976 remote_etag,
977 ..TaskActionState::default()
978 }
979 };
980 let sequence = self.next_task_run.fetch_add(1, Ordering::Relaxed);
981 let run =
982 CacheDigest::blake3(format!("{task}\0{}\0{sequence}", std::process::id()).as_bytes())
983 .hash;
984 self.stats.predictions_loaded.fetch_max(
988 state.predictions.len().try_into().unwrap_or(u64::MAX),
989 Ordering::Relaxed,
990 );
991 let predictions = state.predictions.values().cloned().collect();
992 self.task_actions.lock().unwrap().insert(run.clone(), state);
993 self.spawn_prefetch_predictions(predictions);
994 Ok(run)
995 }
996
997 pub async fn cancel_prefetches(&self) {
999 let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
1000 for task in &tasks {
1001 task.abort();
1002 }
1003 for task in tasks {
1004 if let Err(error) = task.await
1005 && !error.is_cancelled()
1006 {
1007 warn!("remote action prefetch task failed: {error}");
1008 }
1009 }
1010 }
1011
1012 pub async fn wait_for_uploads(&self) {
1019 let Some(uploads) = &self.uploads else {
1020 return;
1021 };
1022 let _timer = AtomicDurationTimer::start(&self.stats.upload_drain_duration_ns);
1023 uploads.drain().await;
1024 }
1025
1026 async fn wait_for_prefetches(&self) {
1027 let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
1028 for task in tasks {
1029 if let Err(error) = task.await {
1030 warn!("remote action prefetch task failed: {error}");
1031 }
1032 }
1033 }
1034
1035 pub async fn commit_task(&self, run: &str) -> Result<()> {
1037 validate_task_identity(run)?;
1038 let state = self
1039 .task_actions
1040 .lock()
1041 .unwrap()
1042 .get(run)
1043 .cloned()
1044 .ok_or_else(|| eyre::eyre!("task action manifest baseline was not loaded"))?;
1045 if !state.baseline_loaded {
1046 bail!("task action manifest baseline was not loaded");
1047 }
1048 let task = state.manifest;
1049 validate_task_identity(&task)?;
1050 let (manifest, introduced) = {
1051 let _write_guard = self.manifest_write_lock.lock().unwrap();
1052 let _file_guard = self.lock_task_manifest(&task)?;
1053 let mut predictions = self
1054 .load_task_manifest(&task)?
1055 .map(|manifest| {
1056 manifest
1057 .predictions
1058 .into_iter()
1059 .map(|prediction| (prediction.invocation.clone(), prediction))
1060 .collect::<BTreeMap<_, _>>()
1061 })
1062 .unwrap_or_default();
1063 let introduced: BTreeSet<CacheDigest> = state
1068 .pending_predictions
1069 .keys()
1070 .filter(|invocation| !predictions.contains_key(*invocation))
1071 .cloned()
1072 .collect();
1073 predictions.extend(state.pending_predictions);
1078 let manifest = TaskActionManifest {
1079 version: TASK_ACTION_MANIFEST_VERSION,
1080 task: task.clone(),
1081 predictions: predictions.into_values().collect(),
1082 };
1083 validate_task_manifest(&manifest, &task)?;
1084 self.persist_task_manifest(&manifest)?;
1085 (manifest, introduced)
1086 };
1087 self.task_actions.lock().unwrap().remove(run);
1088 if self.remote_mode.writes() {
1089 let mut manifest = manifest;
1093 if let Some(uploads) = &self.uploads {
1094 let actions: Vec<CacheDigest> = manifest
1095 .predictions
1096 .iter()
1097 .map(|prediction| prediction.action.clone())
1098 .collect();
1099 let unpublished = uploads.wait_for_actions(&actions).await;
1100 let withheld = manifest
1101 .predictions
1102 .iter()
1103 .filter(|prediction| {
1104 introduced.contains(&prediction.invocation)
1105 && unpublished.contains(&prediction.action)
1106 })
1107 .count();
1108 if withheld > 0 {
1109 warn!(
1112 "{withheld} of {} predicted actions were not published, so the remote task action manifest omits them",
1113 manifest.predictions.len()
1114 );
1115 manifest.predictions.retain(|prediction| {
1116 !(introduced.contains(&prediction.invocation)
1117 && unpublished.contains(&prediction.action))
1118 });
1119 }
1120 }
1121 match self
1122 .put_remote_task_manifest(&task, manifest, state.remote_etag)
1123 .await
1124 {
1125 Ok(remote_manifest) => {
1126 let _write_guard = self.manifest_write_lock.lock().unwrap();
1127 let reconciliation = (|| {
1128 let _file_guard = self.lock_task_manifest(&task)?;
1129 let manifest = match self.load_task_manifest(&task)? {
1130 Some(local) => {
1131 merge_remote_task_manifest(&task, remote_manifest, local).0
1132 }
1133 None => remote_manifest,
1134 };
1135 self.persist_task_manifest(&manifest)
1136 })();
1137 if let Err(error) = reconciliation {
1138 warn!(
1139 "remote task action manifest reconciliation failed for {task}: {error}"
1140 );
1141 }
1142 }
1143 Err(error) => {
1144 self.note_remote_failure();
1145 warn!("remote task action manifest upload failed for {task}: {error}");
1146 }
1147 }
1148 }
1149 Ok(())
1150 }
1151
1152 fn task_manifest_path(&self, task: &str) -> PathBuf {
1153 self.manifest_dir.join(format!("{task}.json"))
1154 }
1155
1156 fn task_manifest_lock_path(&self, task: &str) -> PathBuf {
1157 self.manifest_dir.join("locks").join(format!("{task}.lock"))
1158 }
1159
1160 fn lock_task_manifest(&self, task: &str) -> Result<fslock::LockFile> {
1161 let path = self.task_manifest_lock_path(task);
1162 fs::create_dir_all(path.parent().expect("task manifest lock has a parent"))?;
1163 let mut lock = fslock::LockFile::open(&path)?;
1164 lock.lock()?;
1165 Ok(lock)
1166 }
1167
1168 fn load_task_manifest(&self, task: &str) -> Result<Option<TaskActionManifest>> {
1169 match fs::read(self.task_manifest_path(task)) {
1170 Ok(contents) => Ok(Some(self.parse_task_manifest(task, &contents, false)?)),
1171 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
1172 Err(error) => Err(error.into()),
1173 }
1174 }
1175
1176 fn parse_task_manifest(
1177 &self,
1178 task: &str,
1179 contents: &[u8],
1180 require_canonical: bool,
1181 ) -> Result<TaskActionManifest> {
1182 let manifest: TaskActionManifest = serde_json::from_slice(contents)?;
1183 validate_task_manifest(&manifest, task)?;
1184 if require_canonical && canonical_json(&manifest)? != contents {
1185 bail!("task action manifest is not canonical JSON");
1186 }
1187 Ok(manifest)
1188 }
1189
1190 fn task_manifest_selector(task: &str) -> Result<(Vec<u8>, CacheDigest)> {
1191 TaskActionManifest::selector(task)
1192 }
1193
1194 fn persist_task_manifest(&self, manifest: &TaskActionManifest) -> Result<()> {
1195 let bytes = canonical_json(manifest)?;
1196 fs::create_dir_all(self.manifest_dir.as_path())?;
1197 let mut temporary = tempfile::NamedTempFile::new_in(self.manifest_dir.as_path())?;
1198 std::io::Write::write_all(temporary.as_file_mut(), &bytes)?;
1199 temporary.as_file_mut().sync_all()?;
1200 temporary
1201 .persist(self.task_manifest_path(&manifest.task))
1202 .map_err(|error| error.error)?;
1203 Ok(())
1204 }
1205
1206 async fn get_remote_task_manifest(
1207 &self,
1208 task: &str,
1209 ) -> Result<Option<(TaskActionManifest, String)>> {
1210 let Some(remote) = &self.remote else {
1211 return Ok(None);
1212 };
1213 let (_, selector) = Self::task_manifest_selector(task)?;
1214 let _permit = self.remote_transfers.acquire().await?;
1215 self.stats
1216 .remote_manifest_lookups
1217 .fetch_add(1, Ordering::Relaxed);
1218 let _timer = AtomicDurationTimer::start(&self.stats.remote_manifest_lookup_duration_ns);
1219 let Some(remote_manifest) = remote.get_action_manifest(&selector).await? else {
1220 return Ok(None);
1221 };
1222 let manifest = self.parse_task_manifest(task, &remote_manifest.bytes, true)?;
1223 Ok(Some((manifest, remote_manifest.etag)))
1224 }
1225
1226 async fn put_remote_task_manifest(
1227 &self,
1228 task: &str,
1229 mut manifest: TaskActionManifest,
1230 mut expected_etag: Option<String>,
1231 ) -> Result<TaskActionManifest> {
1232 let Some(remote) = &self.remote else {
1233 return Ok(manifest);
1234 };
1235 let (_, selector) = Self::task_manifest_selector(task)?;
1236 for _ in 0..4 {
1237 let bytes = canonical_json(&manifest)?;
1238 let outcome = {
1239 let _permit = self.remote_transfers.acquire().await?;
1240 remote
1241 .put_action_manifest(&selector, &bytes, expected_etag.as_deref())
1242 .await?
1243 };
1244 match outcome {
1245 ManifestPutOutcome::Stored => return Ok(manifest),
1246 ManifestPutOutcome::PreconditionFailed => {
1247 let Some((remote_manifest, etag)) = self.get_remote_task_manifest(task).await?
1248 else {
1249 expected_etag = None;
1250 continue;
1251 };
1252 manifest = merge_task_manifests(task, Some(remote_manifest), manifest)?;
1253 expected_etag = Some(etag);
1254 }
1255 }
1256 }
1257 bail!("remote task action manifest changed too frequently")
1258 }
1259
1260 pub fn stats(&self) -> AgentStats {
1262 AgentStats {
1263 session_duration_ns: 0,
1264 lookups: self.stats.lookups.load(Ordering::Relaxed),
1265 unconsulted: self.stats.unconsulted.load(Ordering::Relaxed),
1266 hits: self.stats.hits.load(Ordering::Relaxed),
1267 stores: self.stats.stores.load(Ordering::Relaxed),
1268 stored_bytes: self.stats.stored_bytes.load(Ordering::Relaxed),
1269 verifications: self.stats.verifications.load(Ordering::Relaxed),
1270 divergences: self.stats.divergences.load(Ordering::Relaxed),
1271 downloaded_bytes: self.stats.downloaded_bytes.load(Ordering::Relaxed),
1272 uploaded_bytes: self.stats.uploaded_bytes.load(Ordering::Relaxed),
1273 background_uploads: self.stats.background_uploads.load(Ordering::Relaxed),
1274 background_upload_failures: self
1275 .stats
1276 .background_upload_failures
1277 .load(Ordering::Relaxed),
1278 remote_blob_pack_uploads: self.stats.remote_blob_pack_uploads.load(Ordering::Relaxed),
1279 remote_blob_pack_upload_blobs: self
1280 .stats
1281 .remote_blob_pack_upload_blobs
1282 .load(Ordering::Relaxed),
1283 upload_drain_duration_ns: self.stats.upload_drain_duration_ns.load(Ordering::Relaxed),
1284 prefetched_actions: self.stats.prefetched_actions.load(Ordering::Relaxed),
1285 predictions_loaded: self.stats.predictions_loaded.load(Ordering::Relaxed),
1286 bypasses: self.stats.bypasses.lock().unwrap().clone(),
1287 avoided_compiler_duration_ns: self
1288 .stats
1289 .avoided_compiler_duration_ns
1290 .load(Ordering::Relaxed),
1291 compiler: self.stats.compiler.lock().unwrap().clone(),
1292 slow_compilations: self.stats.slow_compilations.lock().unwrap().clone(),
1293 remote_failures: self.stats.remote_failures.load(Ordering::Relaxed),
1294 remote_manifest_lookups: self.stats.remote_manifest_lookups.load(Ordering::Relaxed),
1295 remote_manifest_lookup_duration_ns: self
1296 .stats
1297 .remote_manifest_lookup_duration_ns
1298 .load(Ordering::Relaxed),
1299 remote_action_lookups: self.stats.remote_action_lookups.load(Ordering::Relaxed),
1300 remote_action_lookup_duration_ns: self
1301 .stats
1302 .remote_action_lookup_duration_ns
1303 .load(Ordering::Relaxed),
1304 remote_blob_requests: self.stats.remote_blob_requests.load(Ordering::Relaxed),
1305 remote_blob_pack_requests: self.stats.remote_blob_pack_requests.load(Ordering::Relaxed),
1306 remote_blob_pack_blobs: self.stats.remote_blob_pack_blobs.load(Ordering::Relaxed),
1307 remote_blob_transfer_duration_ns: self
1308 .stats
1309 .remote_blob_transfer_duration_ns
1310 .load(Ordering::Relaxed),
1311 local_cas_write_duration_ns: self
1312 .stats
1313 .local_cas_write_duration_ns
1314 .load(Ordering::Relaxed),
1315 prefetch_runs: self.stats.prefetch_runs.load(Ordering::Relaxed),
1316 prefetch_duration_ns: self.stats.prefetch_duration_ns.load(Ordering::Relaxed),
1317 materialization_duration_ns: self
1318 .stats
1319 .materialization_duration_ns
1320 .load(Ordering::Relaxed),
1321 restored_output_files: self.stats.restored_output_files.load(Ordering::Relaxed),
1322 restored_output_bytes: self.stats.restored_output_bytes.load(Ordering::Relaxed),
1323 reflinked_output_files: self.stats.reflinked_output_files.load(Ordering::Relaxed),
1324 reflinked_output_bytes: self.stats.reflinked_output_bytes.load(Ordering::Relaxed),
1325 copied_output_files: self.stats.copied_output_files.load(Ordering::Relaxed),
1326 copied_output_bytes: self.stats.copied_output_bytes.load(Ordering::Relaxed),
1327 }
1328 }
1329
1330 pub async fn handle_requests(
1337 &self,
1338 requests: impl IntoIterator<Item = AgentRequest>,
1339 ) -> Vec<AgentResponse> {
1340 let mut connection = ConnectionUploads::default();
1341 let mut responses = Vec::new();
1342 for request in requests {
1343 responses.push(self.respond_on(request, &mut connection).await);
1344 }
1345 responses
1346 }
1347
1348 fn write_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1349 Self::digest_lock(&self.write_locks, digest)
1350 }
1351
1352 fn action_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1353 Self::digest_lock(&self.action_locks, digest)
1354 }
1355
1356 fn digest_lock(
1357 locks: &Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>,
1358 digest: &CacheDigest,
1359 ) -> Arc<tokio::sync::Mutex<()>> {
1360 let mut locks = locks.lock().unwrap();
1361 locks.retain(|_, lock| lock.strong_count() > 0);
1362 if let Some(lock) = locks.get(digest).and_then(Weak::upgrade) {
1363 return lock;
1364 }
1365 let lock = Arc::new(tokio::sync::Mutex::new(()));
1366 locks.insert(digest.clone(), Arc::downgrade(&lock));
1367 lock
1368 }
1369
1370 fn spawn_prefetch_predictions(&self, predictions: Vec<ActionPrediction>) {
1371 if predictions.is_empty() || !self.remote_mode.reads() || self.remote.is_none() {
1372 return;
1373 }
1374 let agent = self.clone();
1375 let task = tokio::spawn(async move {
1376 agent.prefetch_predictions(predictions.iter()).await;
1377 });
1378 self.prefetch_tasks.lock().unwrap().push(task);
1379 }
1380
1381 async fn prefetch_predictions<'a>(
1382 &self,
1383 predictions: impl Iterator<Item = &'a ActionPrediction>,
1384 ) {
1385 if !self.remote_mode.reads() || self.remote.is_none() {
1386 return;
1387 }
1388 self.stats.prefetch_runs.fetch_add(1, Ordering::Relaxed);
1389 let _timer = AtomicDurationTimer::start(&self.stats.prefetch_duration_ns);
1390 let mut actions = BTreeMap::new();
1391 for prediction in predictions {
1392 actions
1393 .entry(prediction.action.clone())
1394 .or_insert_with(|| prediction.adapter.clone());
1395 }
1396 match self.prefetch_action_batches(&actions).await {
1399 Ok(true) => return,
1400 Ok(false) => {}
1401 Err(error) => {
1402 self.note_remote_failure();
1403 warn!("remote action batch lookup failed: {error}");
1404 }
1405 }
1406 let mut actions = actions.into_iter();
1407 let mut tasks = tokio::task::JoinSet::new();
1408 for _ in 0..MAX_PREFETCH_TRANSFERS {
1409 let Some((action, adapter)) = actions.next() else {
1410 break;
1411 };
1412 let agent = self.clone();
1413 tasks.spawn(async move { agent.resolve_prefetch_action(action, adapter).await });
1414 }
1415 let mut resolved = Vec::new();
1416 while !tasks.is_empty() {
1417 let result = if resolved.is_empty() {
1418 tasks.join_next().await
1419 } else {
1420 match tokio::time::timeout(PREFETCH_ACTION_BATCH_DELAY, tasks.join_next()).await {
1421 Ok(result) => result,
1422 Err(_) => {
1423 self.prefetch_resolved_actions(std::mem::take(&mut resolved))
1424 .await;
1425 continue;
1426 }
1427 }
1428 };
1429 let Some(result) = result else {
1430 break;
1431 };
1432 match result {
1433 Ok(Ok(Some(action))) => resolved.push(action),
1434 Ok(Ok(None)) => {}
1435 Ok(Err(error)) => {
1436 self.note_remote_failure();
1437 warn!("remote action prefetch failed: {error}");
1438 }
1439 Err(error) => {
1440 self.note_remote_failure();
1441 warn!("remote action prefetch task failed: {error}");
1442 }
1443 }
1444 if let Some((action, adapter)) = actions.next() {
1445 let agent = self.clone();
1446 tasks.spawn(async move { agent.resolve_prefetch_action(action, adapter).await });
1447 }
1448 if resolved.len() == MAX_PREFETCH_ACTION_BATCH {
1449 self.prefetch_resolved_actions(std::mem::take(&mut resolved))
1450 .await;
1451 }
1452 }
1453 if !resolved.is_empty() {
1454 self.prefetch_resolved_actions(resolved).await;
1455 }
1456 }
1457
1458 async fn prefetch_action_batches(
1465 &self,
1466 actions: &BTreeMap<CacheDigest, String>,
1467 ) -> Result<bool> {
1468 let Some(remote) = self.remote.as_deref() else {
1469 return Ok(false);
1470 };
1471 let Some(limit) = remote.action_batch_limit().await? else {
1472 return Ok(false);
1473 };
1474 let wanted: Vec<CacheDigest> = actions
1475 .keys()
1476 .filter(|action| !self.action_is_staged(action))
1477 .cloned()
1478 .collect();
1479 if wanted.is_empty() {
1480 return Ok(true);
1481 }
1482 let mut chunks: Vec<Vec<CacheDigest>> = Vec::new();
1483 for chunk in wanted.chunks(limit) {
1484 chunks.push(chunk.to_vec());
1485 }
1486 let mut lookups = stream::iter(chunks)
1487 .map(|chunk| async move {
1488 let _prefetch_permit = self.prefetch_transfers.acquire().await?;
1489 let _permit = self.remote_transfers.acquire().await?;
1490 self.stats
1491 .remote_action_lookups
1492 .fetch_add(1, Ordering::Relaxed);
1493 let _timer =
1494 AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
1495 remote.get_action_results(&chunk).await
1496 })
1497 .buffer_unordered(MAX_PREFETCH_BATCH_LOOKUPS);
1498 let mut answered = true;
1499 while let Some(lookup) = lookups.next().await {
1500 let results = match lookup {
1501 Ok(Some(results)) => results,
1502 Ok(None) => {
1504 answered = false;
1505 continue;
1506 }
1507 Err(error) => {
1508 self.note_remote_failure();
1509 warn!("remote action batch lookup failed: {error}");
1510 answered = false;
1511 continue;
1512 }
1513 };
1514 let mut resolved = Vec::with_capacity(results.len());
1515 for result in results {
1516 let Some(adapter) = actions.get(&result.action).cloned() else {
1517 continue;
1518 };
1519 let lock = self.action_lock(&result.action);
1520 let _guard = lock.lock().await;
1521 if self.actions.find(&result.action)?.is_some() {
1522 continue;
1523 }
1524 self.pending_remote_actions
1525 .lock()
1526 .unwrap()
1527 .insert(result.action.clone(), result.clone());
1528 resolved.push(PrefetchedAction { adapter, result });
1529 }
1530 while !resolved.is_empty() {
1531 let wave = resolved
1532 .drain(..resolved.len().min(MAX_PREFETCH_ACTION_BATCH))
1533 .collect();
1534 self.prefetch_resolved_actions(wave).await;
1535 }
1536 }
1537 Ok(answered)
1538 }
1539
1540 fn action_is_staged(&self, action: &CacheDigest) -> bool {
1542 if self
1543 .pending_remote_actions
1544 .lock()
1545 .unwrap()
1546 .contains_key(action)
1547 {
1548 return true;
1549 }
1550 self.actions.find(action).is_ok_and(|found| found.is_some())
1551 }
1552
1553 #[cfg(test)]
1554 async fn prefetch_action(&self, action: CacheDigest, adapter: String) -> Result<()> {
1555 if let Some(action) = self.resolve_prefetch_action(action, adapter).await? {
1556 self.prefetch_resolved_actions(vec![action]).await;
1557 }
1558 Ok(())
1559 }
1560
1561 async fn resolve_prefetch_action(
1562 &self,
1563 action: CacheDigest,
1564 adapter: String,
1565 ) -> Result<Option<PrefetchedAction>> {
1566 let remote = self
1567 .remote
1568 .as_ref()
1569 .ok_or_else(|| eyre::eyre!("remote cache is not configured"))?;
1570 let result = {
1571 let lock = self.action_lock(&action);
1572 let _guard = lock.lock().await;
1573 if self.actions.find(&action)?.is_some() {
1574 return Ok(None);
1575 }
1576 if let Some(result) = self
1577 .pending_remote_actions
1578 .lock()
1579 .unwrap()
1580 .get(&action)
1581 .cloned()
1582 {
1583 result
1584 } else {
1585 let _prefetch_permit = self.prefetch_transfers.acquire().await?;
1586 let result = {
1587 let _permit = self.remote_transfers.acquire().await?;
1588 self.get_remote_action_result(remote, &action).await?
1589 };
1590 let Some(result) = result else {
1591 return Ok(None);
1592 };
1593 self.pending_remote_actions
1594 .lock()
1595 .unwrap()
1596 .insert(action.clone(), result.clone());
1597 result
1598 }
1599 };
1600 Ok(Some(PrefetchedAction { adapter, result }))
1601 }
1602
1603 fn prefetch_resolved_actions(&self, actions: Vec<PrefetchedAction>) -> BoxFuture<'_, ()> {
1604 self.prefetch_resolved_actions_inner(actions).boxed()
1605 }
1606
1607 async fn prefetch_resolved_actions_inner(&self, actions: Vec<PrefetchedAction>) {
1608 let Some(remote) = self.remote.as_deref() else {
1609 return;
1610 };
1611 if actions.is_empty() {
1612 return;
1613 }
1614
1615 let mut top_level = BTreeMap::new();
1616 for action in &actions {
1617 for digest in [
1618 Some(&action.result.action),
1619 action.result.metadata.as_ref(),
1620 action.result.output_root.as_ref(),
1621 ]
1622 .into_iter()
1623 .flatten()
1624 {
1625 top_level.insert(digest.clone(), ());
1626 }
1627 }
1628 let mut verified = self
1629 .fetch_remote_blobs(
1630 remote,
1631 top_level.into_keys().collect(),
1632 Some(&self.prefetch_transfers),
1633 )
1634 .await;
1635
1636 let mut next = BTreeMap::new();
1637 let mut pending_directories = BTreeMap::new();
1638 let mut parsed_directories = BTreeMap::new();
1639 let mut rustc_metadata = BTreeMap::new();
1640 for action in &actions {
1641 if action.adapter == "rustc"
1642 && let Some(metadata_digest) = &action.result.metadata
1643 {
1644 match verified
1645 .get(metadata_digest)
1646 .ok_or_else(|| eyre::eyre!("remote rustc action metadata is missing"))
1647 .and_then(|path| Self::parse_rustc_metadata(path))
1648 {
1649 Ok(metadata) => {
1650 queue_prefetch_digest(&verified, &mut next, metadata.stdout.clone());
1651 queue_prefetch_digest(&verified, &mut next, metadata.stderr.clone());
1652 rustc_metadata.insert(metadata_digest.clone(), metadata);
1653 }
1654 Err(error) => warn!(
1655 "remote rustc action metadata prefetch failed for {}: {error}",
1656 action.result.action.hash
1657 ),
1658 }
1659 }
1660 if let Some(output_root) = &action.result.output_root {
1661 pending_directories.insert(output_root.clone(), ());
1662 }
1663 }
1664
1665 let mut seen_directories = BTreeMap::new();
1666 loop {
1667 let mut following = BTreeMap::new();
1668 let mut directory_limit_exceeded = false;
1669 for digest in pending_directories.into_keys() {
1670 following.remove(&digest);
1671 if seen_directories.insert(digest.clone(), ()).is_some() {
1672 continue;
1673 }
1674 if seen_directories.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
1675 warn!("remote action output tree is too large to prefetch");
1676 following.clear();
1677 break;
1678 }
1679 match verified
1680 .get(&digest)
1681 .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))
1682 .and_then(|path| Self::parse_cache_directory(path))
1683 {
1684 Ok(directory) => {
1685 for file in &directory.files {
1686 queue_prefetch_digest(&verified, &mut next, file.digest.clone());
1687 if next.len() >= MAX_PREFETCH_OBJECTS_PER_WAVE {
1688 self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1689 .await;
1690 }
1691 }
1692 for child in &directory.directories {
1693 if !queue_prefetch_directory(
1694 &seen_directories,
1695 &mut following,
1696 child.digest.clone(),
1697 MAX_PREFETCH_DIRECTORY_OBJECTS,
1698 ) {
1699 warn!("remote action output tree is too large to prefetch");
1700 directory_limit_exceeded = true;
1701 break;
1702 }
1703 queue_prefetch_digest(&verified, &mut next, child.digest.clone());
1704 if next.len() >= MAX_PREFETCH_OBJECTS_PER_WAVE {
1705 self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1706 .await;
1707 }
1708 }
1709 parsed_directories.insert(digest, directory);
1710 }
1711 Err(error) => warn!(
1712 "remote action output directory prefetch failed for {}: {error}",
1713 digest.hash
1714 ),
1715 }
1716 if directory_limit_exceeded {
1717 following.clear();
1718 break;
1719 }
1720 }
1721 self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1722 .await;
1723 if following.is_empty() {
1724 break;
1725 }
1726 pending_directories = following;
1727 }
1728
1729 for action in actions {
1730 match Self::validate_prefetched_action(
1731 &action,
1732 &verified,
1733 &rustc_metadata,
1734 &parsed_directories,
1735 ) {
1736 Ok(()) => {
1737 if let Err(error) = self.actions.store(&action.result) {
1738 warn!(
1739 "remote action prefetch could not publish {}: {error}",
1740 action.result.action.hash
1741 );
1742 continue;
1743 }
1744 self.pending_remote_actions
1745 .lock()
1746 .unwrap()
1747 .remove(&action.result.action);
1748 self.stats
1749 .prefetched_actions
1750 .fetch_add(1, Ordering::Relaxed);
1751 }
1752 Err(error) => warn!(
1753 "remote action prefetch was incomplete for {}: {error}",
1754 action.result.action.hash
1755 ),
1756 }
1757 }
1758 }
1759
1760 async fn flush_prefetch_digest_batch(
1761 &self,
1762 remote: &RemoteCacheClient,
1763 verified: &mut BTreeMap<CacheDigest, PathBuf>,
1764 pending: &mut BTreeMap<CacheDigest, ()>,
1765 ) {
1766 if pending.is_empty() {
1767 return;
1768 }
1769 let digests = std::mem::take(pending).into_keys().collect();
1770 verified.extend(
1771 self.fetch_remote_blobs(remote, digests, Some(&self.prefetch_transfers))
1772 .await,
1773 );
1774 }
1775
1776 async fn fetch_remote_blobs(
1777 &self,
1778 remote: &RemoteCacheClient,
1779 digests: Vec<CacheDigest>,
1780 prefetch_limit: Option<&tokio::sync::Semaphore>,
1781 ) -> BTreeMap<CacheDigest, PathBuf> {
1782 let mut verified = BTreeMap::new();
1783 let mut missing = BTreeMap::new();
1784 for digest in digests {
1785 match self.find_verified_blob(&digest) {
1786 Ok(Some(path)) => {
1787 verified.insert(digest, path);
1788 }
1789 Ok(None) => {
1790 missing.insert(digest, ());
1791 }
1792 Err(error) => warn!(
1793 "local cache blob lookup failed for {}: {error}",
1794 digest.hash
1795 ),
1796 }
1797 }
1798 if missing.is_empty() {
1799 return verified;
1800 }
1801
1802 let mut pack_candidates = missing.clone();
1803 while !pack_candidates.is_empty() {
1804 let candidates = match blob_pack_chunk(
1805 &pack_candidates.keys().cloned().collect::<Vec<_>>(),
1806 BlobPackLimits {
1807 max_items: MAX_STAGED_BLOB_PACK_ITEMS,
1808 max_bytes: MAX_STAGED_BLOB_PACK_BYTES,
1809 },
1810 ) {
1811 Ok(candidates) if !candidates.is_empty() => candidates,
1812 Ok(_) => break,
1813 Err(error) => {
1814 warn!("remote cache blob pack skipped: {error}");
1815 break;
1816 }
1817 };
1818 let mut pack_guards = BTreeMap::new();
1822 for digest in candidates {
1823 let guard = self.write_lock(&digest).lock_owned().await;
1824 match self.find_verified_blob(&digest) {
1825 Ok(Some(path)) => {
1826 pack_candidates.remove(&digest);
1827 missing.remove(&digest);
1828 verified.insert(digest, path);
1829 }
1830 Ok(None) => {
1831 pack_guards.insert(digest, guard);
1832 }
1833 Err(error) => {
1834 warn!(
1835 "local cache blob lookup failed for {}: {error}",
1836 digest.hash
1837 );
1838 pack_guards.insert(digest, guard);
1839 }
1840 }
1841 }
1842 let requested = pack_guards.keys().cloned().collect::<Vec<_>>();
1843 if requested.is_empty() {
1844 continue;
1845 }
1846 let requested_bytes = requested
1847 .iter()
1848 .fold(0_u64, |total, digest| total.saturating_add(digest.size));
1849 let pack_reservation = match self
1850 .reserve_remote_download_up_to(requested_bytes.min(MAX_STAGED_BLOB_PACK_BYTES))
1851 {
1852 Ok(reservation) if reservation.bytes() > 0 => reservation,
1853 Ok(_) => break,
1854 Err(error) => {
1855 warn!("remote cache blob pack skipped: {error}");
1856 break;
1857 }
1858 };
1859 let (pack, transfer_duration_ns) = {
1860 let _prefetch_permit = match prefetch_limit {
1861 Some(limit) => match limit.acquire().await {
1862 Ok(permit) => Some(permit),
1863 Err(error) => {
1864 warn!(
1865 "remote cache blob pack could not acquire prefetch limit: {error}"
1866 );
1867 break;
1868 }
1869 },
1870 None => None,
1871 };
1872 let _transfer_permit = match self.remote_transfers.acquire().await {
1873 Ok(permit) => permit,
1874 Err(error) => {
1875 warn!("remote cache blob pack could not acquire transfer limit: {error}");
1876 break;
1877 }
1878 };
1879 let transfer_started = Instant::now();
1880 let pack = remote
1881 .get_blob_pack_with_limit(
1882 &requested,
1883 self.remote_staging_dir.as_path(),
1884 pack_reservation.bytes(),
1885 )
1886 .await;
1887 (pack, duration_ns(transfer_started))
1888 };
1889 let pack = match pack {
1890 Ok(Some(pack)) => pack,
1891 Ok(None) => break,
1892 Err(error) => {
1893 atomic_saturating_add(
1894 &self.stats.remote_blob_transfer_duration_ns,
1895 transfer_duration_ns,
1896 );
1897 warn!(
1898 "remote cache blob pack failed; falling back to individual blobs: {error}"
1899 );
1900 break;
1901 }
1902 };
1903 pack_reservation.commit(pack.payload_bytes);
1904 atomic_saturating_add(
1905 &self.stats.remote_blob_transfer_duration_ns,
1906 transfer_duration_ns,
1907 );
1908 atomic_saturating_add(&self.stats.remote_blob_pack_requests, pack.requests);
1909 atomic_saturating_add(&self.stats.remote_blob_pack_blobs, pack.blob_count);
1910 atomic_saturating_add(&self.stats.downloaded_bytes, pack.payload_bytes);
1911 if pack.requested.is_empty() {
1912 for digest in &requested {
1916 pack_candidates.remove(digest);
1917 }
1918 continue;
1919 }
1920 for digest in &pack.requested {
1921 pack_candidates.remove(digest);
1922 }
1923 let mut ingests = stream::iter(pack.blobs.into_iter().map(|(digest, source)| {
1924 let digest_for_result = digest.clone();
1925 let guard = pack_guards
1926 .remove(&digest)
1927 .expect("requested packed blob has a write lock");
1928 async move {
1929 (
1930 digest_for_result,
1931 self.ingest_packed_blob(digest, source, guard).await,
1932 )
1933 }
1934 }))
1935 .buffer_unordered(MAX_PREFETCH_TRANSFERS);
1936 while let Some((digest, result)) = ingests.next().await {
1937 match result {
1938 Ok(path) => {
1939 missing.remove(&digest);
1940 verified.insert(digest, path);
1941 }
1942 Err(error) => warn!(
1943 "remote cache packed blob ingest failed for {}: {error}",
1944 digest.hash
1945 ),
1946 }
1947 }
1948 }
1949
1950 let mut transfers = stream::iter(missing.into_keys().map(|digest| {
1951 let digest_for_result = digest.clone();
1952 async move {
1953 (
1954 digest_for_result,
1955 self.fetch_remote_blob_with_limit(remote, &digest, prefetch_limit)
1956 .await,
1957 )
1958 }
1959 }))
1960 .buffer_unordered(MAX_PREFETCH_TRANSFERS);
1961 while let Some((digest, result)) = transfers.next().await {
1962 match result {
1963 Ok(path) => {
1964 verified.insert(digest, path);
1965 }
1966 Err(error) => warn!(
1967 "remote cache blob prefetch failed for {}: {error}",
1968 digest.hash
1969 ),
1970 }
1971 }
1972 verified
1973 }
1974
1975 async fn ingest_packed_blob(
1976 &self,
1977 digest: CacheDigest,
1978 source: PathBuf,
1979 _guard: tokio::sync::OwnedMutexGuard<()>,
1980 ) -> Result<PathBuf> {
1981 let digest_size = digest.size;
1982 let agent = self.clone();
1983 let (path, stored, cas_duration_ns) = tokio::task::spawn_blocking(move || {
1984 if let Some(path) = agent.find_verified_blob(&digest)? {
1985 return Ok::<_, eyre::Report>((path, false, 0));
1986 }
1987 let cas_started = Instant::now();
1988 let path = agent.cas.store_verified_file(&digest, &source)?;
1989 let cas_duration_ns = duration_ns(cas_started);
1990 agent.remember_verified_blob(&digest, &path);
1991 Ok((path, true, cas_duration_ns))
1992 })
1993 .await??;
1994 atomic_saturating_add(&self.stats.local_cas_write_duration_ns, cas_duration_ns);
1995 if stored {
1996 self.stats.stores.fetch_add(1, Ordering::Relaxed);
1997 atomic_saturating_add(&self.stats.stored_bytes, digest_size);
1998 }
1999 Ok(path)
2000 }
2001
2002 fn parse_rustc_metadata(path: &Path) -> Result<RustcMetadata> {
2003 let bytes = fs::read(path)?;
2004 let metadata: RustcMetadata = serde_json::from_slice(&bytes)?;
2005 if metadata.version != 1 || metadata.kind != "rustc" || canonical_json(&metadata)? != bytes
2006 {
2007 bail!("remote rustc action metadata is invalid");
2008 }
2009 Ok(metadata)
2010 }
2011
2012 fn parse_cache_directory(path: &Path) -> Result<CacheDirectory> {
2013 let bytes = fs::read(path)?;
2014 let directory: CacheDirectory = serde_json::from_slice(&bytes)?;
2015 if directory.version != 1 || canonical_json(&directory)? != bytes {
2016 bail!("remote action output directory is invalid");
2017 }
2018 Ok(directory)
2019 }
2020
2021 #[cfg(test)]
2022 fn load_cache_directory(&self, digest: &CacheDigest) -> Result<CacheDirectory> {
2023 let path = self
2024 .find_verified_blob(digest)?
2025 .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))?;
2026 Self::parse_cache_directory(&path)
2027 }
2028
2029 fn validate_prefetched_action(
2030 action: &PrefetchedAction,
2031 verified: &BTreeMap<CacheDigest, PathBuf>,
2032 rustc_metadata: &BTreeMap<CacheDigest, RustcMetadata>,
2033 directories: &BTreeMap<CacheDigest, CacheDirectory>,
2034 ) -> Result<()> {
2035 if !verified.contains_key(&action.result.action) {
2036 bail!("remote action descriptor is missing");
2037 }
2038 if let Some(metadata) = &action.result.metadata {
2039 if action.adapter == "rustc" {
2040 let metadata = rustc_metadata
2041 .get(metadata)
2042 .ok_or_else(|| eyre::eyre!("remote rustc action metadata is missing"))?;
2043 for digest in [&metadata.stdout, &metadata.stderr] {
2044 if !verified.contains_key(digest) {
2045 bail!("remote rustc action diagnostic blob is missing");
2046 }
2047 }
2048 } else if !verified.contains_key(metadata) {
2049 bail!("remote action metadata is missing");
2050 }
2051 }
2052 let mut pending = action
2053 .result
2054 .output_root
2055 .iter()
2056 .cloned()
2057 .collect::<Vec<_>>();
2058 let mut seen = BTreeMap::new();
2059 while let Some(digest) = pending.pop() {
2060 if seen.insert(digest.clone(), ()).is_some() {
2061 continue;
2062 }
2063 if seen.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
2064 bail!("remote action output tree is too large");
2065 }
2066 let directory = directories
2067 .get(&digest)
2068 .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))?;
2069 for file in &directory.files {
2070 if !verified.contains_key(&file.digest) {
2071 bail!("remote action output file is missing");
2072 }
2073 }
2074 pending.extend(
2075 directory
2076 .directories
2077 .iter()
2078 .map(|directory| directory.digest.clone()),
2079 );
2080 }
2081 Ok(())
2082 }
2083
2084 #[cfg(test)]
2085 async fn prefetch_output_tree(
2086 &self,
2087 remote: &RemoteCacheClient,
2088 output_root: &CacheDigest,
2089 ) -> Result<()> {
2090 let mut pending = vec![output_root.clone()];
2091 let mut seen = BTreeMap::new();
2092 while let Some(digest) = pending.pop() {
2093 if seen.insert(digest.clone(), ()).is_some() {
2094 continue;
2095 }
2096 if seen.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
2097 bail!("remote action output tree is too large");
2098 }
2099 self.fetch_remote_blob_with_limit(remote, &digest, Some(&self.prefetch_transfers))
2100 .await?;
2101 let directory = self.load_cache_directory(&digest)?;
2102 let mut transfers = stream::iter(directory.files.into_iter().map(|file| async move {
2103 self.fetch_remote_blob_with_limit(
2104 remote,
2105 &file.digest,
2106 Some(&self.prefetch_transfers),
2107 )
2108 .await
2109 .map(|_| ())
2110 }))
2111 .buffer_unordered(MAX_PREFETCH_TRANSFERS);
2112 while let Some(result) = transfers.next().await {
2113 result?;
2114 }
2115 pending.extend(
2116 directory
2117 .directories
2118 .into_iter()
2119 .map(|directory| directory.digest),
2120 );
2121 }
2122 Ok(())
2123 }
2124
2125 async fn fetch_remote_blob(
2126 &self,
2127 remote: &RemoteCacheClient,
2128 digest: &CacheDigest,
2129 ) -> Result<PathBuf> {
2130 self.fetch_remote_blob_with_limit(remote, digest, None)
2131 .await
2132 }
2133
2134 async fn fetch_remote_blob_with_limit(
2135 &self,
2136 remote: &RemoteCacheClient,
2137 digest: &CacheDigest,
2138 prefetch_limit: Option<&tokio::sync::Semaphore>,
2139 ) -> Result<PathBuf> {
2140 let lock = self.write_lock(digest);
2141 let _guard = lock.lock().await;
2142 if let Some(path) = self.find_verified_blob(digest)? {
2143 return Ok(path);
2144 }
2145 let _prefetch_permit = match prefetch_limit {
2146 Some(limit) => Some(limit.acquire().await?),
2147 None => None,
2148 };
2149 let _permit = self.remote_transfers.acquire().await?;
2150 let reservation = self.reserve_remote_download(digest.size)?;
2151 self.stats
2152 .remote_blob_requests
2153 .fetch_add(1, Ordering::Relaxed);
2154 let transfer_timer =
2155 AtomicDurationTimer::start(&self.stats.remote_blob_transfer_duration_ns);
2156 let temporary = remote
2157 .get_blob_file(digest, self.remote_staging_dir.as_path())
2158 .await?;
2159 drop(transfer_timer);
2160 let _cas_timer = AtomicDurationTimer::start(&self.stats.local_cas_write_duration_ns);
2161 let path = self.cas.store_verified_file(digest, temporary.path())?;
2162 reservation.commit(digest.size);
2163 self.remember_verified_blob(digest, &path);
2164 self.stats.stores.fetch_add(1, Ordering::Relaxed);
2165 self.stats
2166 .stored_bytes
2167 .fetch_add(digest.size, Ordering::Relaxed);
2168 self.stats
2169 .downloaded_bytes
2170 .fetch_add(digest.size, Ordering::Relaxed);
2171 Ok(path)
2172 }
2173
2174 #[cfg(test)]
2179 async fn respond(&self, request: AgentRequest) -> AgentResponse {
2180 self.respond_on(request, &mut ConnectionUploads::default())
2181 .await
2182 }
2183
2184 async fn respond_on(
2185 &self,
2186 request: AgentRequest,
2187 connection: &mut ConnectionUploads,
2188 ) -> AgentResponse {
2189 let result = match request {
2190 AgentRequest::BeginTask { task } => self
2191 .begin_task(&task)
2192 .await
2193 .map(|run| AgentResponse::TaskBegun { run }),
2194 AgentRequest::CommitTask { run } => self
2195 .commit_task(&run)
2196 .await
2197 .map(|()| AgentResponse::TaskCommitted),
2198 AgentRequest::FindBlob { digest } => self.find_blob(&digest).await,
2199 AgentRequest::FindBlobs { digests } => self.find_blobs(digests).await,
2200 AgentRequest::StoreBlob { digest, source } => {
2201 self.store_blob(&digest, &source, connection).await
2202 }
2203 AgentRequest::FindActionResult { action } => {
2204 self.stats.lookups.fetch_add(1, Ordering::Relaxed);
2205 self.find_action_result(&action).await
2206 }
2207 AgentRequest::RecordActionHit {
2208 action,
2209 restore,
2210 crate_name,
2211 } => self.record_action_hit(&action, restore, crate_name),
2212 AgentRequest::RecordBypass { kind } => {
2213 *self
2214 .stats
2215 .bypasses
2216 .lock()
2217 .unwrap()
2218 .entry(kind.clone())
2219 .or_insert(0) += 1;
2220 self.emit(|| AgentEvent::Bypass { kind });
2221 Ok(AgentResponse::BypassRecorded)
2222 }
2223 AgentRequest::RecordUnconsulted => {
2224 self.stats.unconsulted.fetch_add(1, Ordering::Relaxed);
2225 self.emit(|| AgentEvent::Unconsulted);
2226 Ok(AgentResponse::UnconsultedRecorded)
2227 }
2228 AgentRequest::RecordWarning { message } => self.record_warning(message),
2229 AgentRequest::RecordCompilerInvocation {
2230 outcome,
2231 crate_name,
2232 duration_ns,
2233 } => self.record_compiler_invocation(&outcome, crate_name.as_deref(), duration_ns),
2234 AgentRequest::RecordActionVerification { matched, restore } => {
2235 self.record_materialization(restore);
2236 self.stats.verifications.fetch_add(1, Ordering::Relaxed);
2237 if !matched {
2238 self.stats.divergences.fetch_add(1, Ordering::Relaxed);
2239 }
2240 self.emit(|| AgentEvent::Verification { matched, restore });
2241 Ok(AgentResponse::ActionVerificationRecorded)
2242 }
2243 AgentRequest::StoreActionResult { result } => {
2244 self.store_action_result(&result, connection).await
2245 }
2246 AgentRequest::FindActionPrediction { task, invocation } => {
2247 self.find_action_prediction(&task, &invocation)
2248 }
2249 AgentRequest::RecordActionPrediction { task, prediction } => {
2250 self.record_action_prediction(&task, prediction)
2251 }
2252 AgentRequest::FindExecutableIdentity {
2253 executable,
2254 environment,
2255 } => self.find_executable_identity(executable, environment),
2256 AgentRequest::StoreExecutableIdentity {
2257 executable,
2258 environment,
2259 stdout,
2260 } => self.store_executable_identity(executable, environment, stdout),
2261 AgentRequest::Hello { .. } => {
2262 Err(eyre::eyre!("hello is only valid as the first request"))
2263 }
2264 };
2265 result.unwrap_or_else(|error| AgentResponse::Error {
2266 message: error.to_string(),
2267 })
2268 }
2269
2270 async fn find_blob(&self, digest: &CacheDigest) -> Result<AgentResponse> {
2271 if let Some(path) = self.find_verified_blob(digest)? {
2272 return Ok(AgentResponse::Blob { path: Some(path) });
2273 }
2274 if !self.remote_mode.reads() {
2275 return Ok(AgentResponse::Blob { path: None });
2276 }
2277 let Some(remote) = &self.remote else {
2278 return Ok(AgentResponse::Blob { path: None });
2279 };
2280 match self.fetch_remote_blob(remote, digest).await {
2281 Ok(path) => Ok(AgentResponse::Blob { path: Some(path) }),
2282 Err(error) => {
2283 warn!(
2284 "remote cache blob lookup failed for {}: {error}",
2285 digest.hash
2286 );
2287 Ok(AgentResponse::Blob { path: None })
2288 }
2289 }
2290 }
2291
2292 async fn find_blobs(&self, digests: Vec<CacheDigest>) -> Result<AgentResponse> {
2293 let mut paths = BTreeMap::new();
2294 let mut missing = Vec::new();
2295 for digest in &digests {
2296 match self.find_verified_blob(digest)? {
2297 Some(path) => {
2298 paths.insert(digest.clone(), path);
2299 }
2300 None => {
2301 missing.push(digest.clone());
2302 }
2303 }
2304 }
2305
2306 if !missing.is_empty()
2307 && self.remote_mode.reads()
2308 && let Some(remote) = &self.remote
2309 {
2310 paths.extend(self.fetch_remote_blobs(remote, missing, None).await);
2311 }
2312
2313 Ok(AgentResponse::Blobs {
2314 paths: digests
2315 .into_iter()
2316 .map(|digest| paths.get(&digest).cloned())
2317 .collect(),
2318 })
2319 }
2320
2321 async fn store_blob(
2322 &self,
2323 digest: &CacheDigest,
2324 source: &Path,
2325 connection: &mut ConnectionUploads,
2326 ) -> Result<AgentResponse> {
2327 let path = {
2328 let lock = self.write_lock(digest);
2329 let _guard = lock.lock().await;
2330 if let Some(path) = self.find_verified_blob(digest)? {
2331 path
2332 } else {
2333 let path = self.cas.store_file(digest, source)?;
2334 self.remember_verified_blob(digest, &path);
2335 self.stats.stores.fetch_add(1, Ordering::Relaxed);
2336 self.stats
2337 .stored_bytes
2338 .fetch_add(digest.size, Ordering::Relaxed);
2339 path
2340 }
2341 };
2342 if let Some(uploads) = &self.uploads {
2345 uploads.queue_blob(digest, path.clone(), connection);
2346 }
2347 Ok(AgentResponse::Stored { path })
2348 }
2349
2350 fn find_verified_blob(&self, digest: &CacheDigest) -> Result<Option<PathBuf>> {
2351 let remembered = self.verified_blobs.lock().unwrap().get(digest).cloned();
2352 if let Some(remembered) = remembered {
2353 if remembered.is_unchanged() {
2364 return Ok(Some(remembered.path));
2365 }
2366 self.verified_blobs.lock().unwrap().remove(digest);
2367 }
2368 let path = self.cas.find(digest)?;
2369 if let Some(path) = &path {
2370 self.remember_verified_blob(digest, path);
2371 }
2372 Ok(path)
2373 }
2374
2375 fn remember_verified_blob(&self, digest: &CacheDigest, path: &Path) {
2376 let Some(verified) = VerifiedBlob::describe(path) else {
2377 return;
2381 };
2382 self.verified_blobs
2383 .lock()
2384 .unwrap()
2385 .insert(digest.clone(), verified);
2386 }
2387
2388 async fn find_action_result(&self, action: &CacheDigest) -> Result<AgentResponse> {
2389 if let Some(result) = self.actions.find(action)? {
2390 return Ok(AgentResponse::ActionResult {
2391 result: Some(result),
2392 });
2393 }
2394 if !self.remote_mode.reads() {
2395 return Ok(AgentResponse::ActionResult { result: None });
2396 }
2397 let Some(remote) = &self.remote else {
2398 return Ok(AgentResponse::ActionResult { result: None });
2399 };
2400 let lock = self.action_lock(action);
2401 let _guard = lock.lock().await;
2402 if let Some(result) = self.actions.find(action)? {
2403 return Ok(AgentResponse::ActionResult {
2404 result: Some(result),
2405 });
2406 }
2407 if let Some(result) = self
2408 .pending_remote_actions
2409 .lock()
2410 .unwrap()
2411 .get(action)
2412 .cloned()
2413 {
2414 return Ok(AgentResponse::ActionResult {
2415 result: Some(result),
2416 });
2417 }
2418 let _permit = self.remote_transfers.acquire().await?;
2419 match self.get_remote_action_result(remote, action).await {
2420 Ok(Some(result)) => {
2421 self.pending_remote_actions
2422 .lock()
2423 .unwrap()
2424 .insert(action.clone(), result.clone());
2425 Ok(AgentResponse::ActionResult {
2426 result: Some(result),
2427 })
2428 }
2429 Ok(None) => Ok(AgentResponse::ActionResult { result: None }),
2430 Err(error) => {
2431 self.note_remote_failure();
2432 warn!(
2433 "remote cache action lookup failed for {}: {error}",
2434 action.hash
2435 );
2436 Ok(AgentResponse::ActionResult { result: None })
2437 }
2438 }
2439 }
2440
2441 async fn store_action_result(
2442 &self,
2443 result: &RemoteActionResult,
2444 connection: &ConnectionUploads,
2445 ) -> Result<AgentResponse> {
2446 let path = self.actions.store(result)?;
2447 if let Some(uploads) = &self.uploads {
2448 uploads.queue_action_result(result, connection);
2449 }
2450 Ok(AgentResponse::ActionStored { path })
2451 }
2452
2453 fn note_remote_failure(&self) {
2455 self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
2456 }
2457
2458 async fn get_remote_action_result(
2459 &self,
2460 remote: &RemoteCacheClient,
2461 action: &CacheDigest,
2462 ) -> Result<Option<RemoteActionResult>> {
2463 self.stats
2464 .remote_action_lookups
2465 .fetch_add(1, Ordering::Relaxed);
2466 let _timer = AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
2467 remote.get_action_result(action).await
2468 }
2469
2470 fn record_action_hit(
2471 &self,
2472 action: &CacheDigest,
2473 restore: RestoreStats,
2474 crate_name: Option<String>,
2475 ) -> Result<AgentResponse> {
2476 validate_crate_name(crate_name.as_deref())?;
2477 if self.actions.find(action)?.is_none() {
2478 let pending = self.pending_remote_actions.lock().unwrap().remove(action);
2479 if let Some(result) = pending {
2480 self.actions.store(&result)?;
2481 } else {
2482 bail!("cannot record a hit for a missing action result");
2483 }
2484 }
2485 self.record_restore(restore);
2486 self.stats.hits.fetch_add(1, Ordering::Relaxed);
2487 self.emit(|| AgentEvent::ActionHit {
2488 crate_name,
2489 restore,
2490 });
2491 Ok(AgentResponse::ActionHitRecorded)
2492 }
2493
2494 fn record_restore(&self, restore: RestoreStats) {
2495 self.record_materialization(restore);
2496 atomic_saturating_add(
2497 &self.stats.avoided_compiler_duration_ns,
2498 restore.avoided_compiler_duration_ns,
2499 );
2500 atomic_saturating_add(&self.stats.restored_output_files, restore.output_files);
2501 atomic_saturating_add(&self.stats.restored_output_bytes, restore.output_bytes);
2502 atomic_saturating_add(
2503 &self.stats.reflinked_output_files,
2504 restore.reflinked_output_files,
2505 );
2506 atomic_saturating_add(
2507 &self.stats.reflinked_output_bytes,
2508 restore.reflinked_output_bytes,
2509 );
2510 atomic_saturating_add(&self.stats.copied_output_files, restore.copied_output_files);
2511 atomic_saturating_add(&self.stats.copied_output_bytes, restore.copied_output_bytes);
2512 }
2513
2514 fn record_compiler_invocation(
2515 &self,
2516 outcome: &str,
2517 crate_name: Option<&str>,
2518 duration_ns: u64,
2519 ) -> Result<AgentResponse> {
2520 if !matches!(
2521 outcome,
2522 "miss" | "unconsulted" | "bypass" | "verification" | "incremental"
2523 ) {
2524 bail!("invalid compiler invocation outcome");
2525 }
2526 validate_crate_name(crate_name)?;
2527 let mut compiler = self.stats.compiler.lock().unwrap();
2528 let stats = compiler.entry(outcome.to_string()).or_default();
2529 stats.invocations = stats.invocations.saturating_add(1);
2530 stats.duration_ns = stats.duration_ns.saturating_add(duration_ns);
2531 drop(compiler);
2532 if outcome != "verification"
2533 && let Some(crate_name) = crate_name.filter(|name| !name.is_empty())
2534 {
2535 let mut slow = self.stats.slow_compilations.lock().unwrap();
2536 let duration = slow.entry(crate_name.to_string()).or_default();
2537 *duration = duration.saturating_add(duration_ns);
2538 }
2539 self.emit(|| AgentEvent::CompilerInvocation {
2540 outcome: outcome.to_string(),
2541 crate_name: crate_name.map(str::to_string),
2542 duration_ns,
2543 });
2544 Ok(AgentResponse::CompilerInvocationRecorded)
2545 }
2546
2547 fn record_materialization(&self, restore: RestoreStats) {
2548 atomic_saturating_add(&self.stats.materialization_duration_ns, restore.duration_ns);
2549 }
2550
2551 fn record_warning(&self, message: String) -> Result<AgentResponse> {
2560 if message.is_empty()
2561 || message.len() > MAX_WARNING_BYTES
2562 || message.contains(['\n', '\r', '\0'])
2563 {
2564 bail!("invalid shim warning");
2565 }
2566 let mut warnings = self.warnings.lock().unwrap();
2567 if !warnings.contains(&message) && warnings.len() < MAX_WARNINGS {
2568 eprintln!("mbx[warning]: {message}");
2569 self.emit(|| AgentEvent::Warning {
2570 message: message.clone(),
2571 });
2572 warnings.insert(message);
2573 }
2574 Ok(AgentResponse::WarningRecorded)
2575 }
2576
2577 fn find_action_prediction(
2578 &self,
2579 task: &str,
2580 invocation: &CacheDigest,
2581 ) -> Result<AgentResponse> {
2582 validate_task_identity(task)?;
2583 invocation.validate()?;
2584 let prediction = self
2585 .task_actions
2586 .lock()
2587 .unwrap()
2588 .get(task)
2589 .and_then(|state| state.predictions.get(invocation))
2590 .cloned();
2591 Ok(AgentResponse::ActionPrediction { prediction })
2592 }
2593
2594 fn record_action_prediction(
2595 &self,
2596 task: &str,
2597 prediction: ActionPrediction,
2598 ) -> Result<AgentResponse> {
2599 validate_task_identity(task)?;
2600 validate_action_prediction(&prediction)?;
2601 let mut tasks = self.task_actions.lock().unwrap();
2602 let state = tasks.entry(task.to_string()).or_default();
2603 if !state.predictions.contains_key(&prediction.invocation)
2604 && state.predictions.len() >= MAX_TASK_ACTION_PREDICTIONS
2605 {
2606 bail!("task action manifest contains too many predictions");
2607 }
2608 state
2609 .predictions
2610 .insert(prediction.invocation.clone(), prediction.clone());
2611 state
2612 .pending_predictions
2613 .insert(prediction.invocation.clone(), prediction);
2614 Ok(AgentResponse::ActionPredictionRecorded)
2615 }
2616
2617 fn executable_identity_key(
2618 &self,
2619 executable: PathBuf,
2620 environment: BTreeMap<String, Option<String>>,
2621 ) -> Result<ExecutableIdentityKey> {
2622 if !environment.keys().all(|name| {
2627 matches!(
2628 name.as_str(),
2629 "RUSTUP_HOME" | "RUSTUP_TOOLCHAIN" | "SDKROOT" | "MACOSX_DEPLOYMENT_TARGET"
2630 )
2631 }) {
2632 bail!("executable identity contains an unsupported environment variable");
2633 }
2634 Ok(ExecutableIdentityKey {
2635 executable,
2636 environment,
2637 })
2638 }
2639
2640 fn find_executable_identity(
2641 &self,
2642 executable: PathBuf,
2643 environment: BTreeMap<String, Option<String>>,
2644 ) -> Result<AgentResponse> {
2645 let key = self.executable_identity_key(executable, environment)?;
2646 let stdout = self
2647 .executable_identities
2648 .lock()
2649 .unwrap()
2650 .get(&key)
2651 .cloned();
2652 Ok(AgentResponse::ExecutableIdentity { stdout })
2653 }
2654
2655 fn store_executable_identity(
2656 &self,
2657 executable: PathBuf,
2658 environment: BTreeMap<String, Option<String>>,
2659 stdout: Vec<u8>,
2660 ) -> Result<AgentResponse> {
2661 if stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE {
2662 bail!("executable identity exceeds {MAX_EXECUTABLE_IDENTITY_SIZE} bytes");
2663 }
2664 let key = self.executable_identity_key(executable, environment)?;
2665 let mut identities = self.executable_identities.lock().unwrap();
2666 let is_new = !identities.contains_key(&key);
2667 let previous_size = identities.get(&key).map_or(0, Vec::len);
2668 if is_new && identities.len() >= MAX_EXECUTABLE_IDENTITIES {
2669 bail!("executable identity cache contains too many entries");
2670 }
2671 let retained_bytes = identities.values().map(Vec::len).sum::<usize>();
2672 if retained_bytes - previous_size + stdout.len() > MAX_EXECUTABLE_IDENTITY_BYTES {
2673 bail!("executable identity cache contains too many bytes");
2674 }
2675 identities.insert(key, stdout.clone());
2676 Ok(AgentResponse::ExecutableIdentity {
2677 stdout: Some(stdout),
2678 })
2679 }
2680
2681 pub async fn handle_connection<S>(&self, stream: S) -> Result<()>
2683 where
2684 S: AsyncRead + AsyncWrite + Unpin,
2685 {
2686 let (reader, mut writer) = tokio::io::split(stream);
2687 let mut reader = BufReader::new(reader);
2688 let hello = read_request(&mut reader)
2689 .await?
2690 .ok_or_else(|| eyre::eyre!("connection closed before the agent handshake"))?;
2691 let request: AgentRequest = serde_json::from_str(&hello)?;
2692 match request {
2693 AgentRequest::Hello {
2694 protocol,
2695 client_version,
2696 } if protocol == AGENT_PROTOCOL_VERSION && client_version == self.version.as_ref() => {}
2697 AgentRequest::Hello { protocol, .. } if protocol != AGENT_PROTOCOL_VERSION => {
2698 send_response(
2699 &mut writer,
2700 &AgentResponse::Error {
2701 message: format!(
2702 "unsupported agent protocol {protocol}; expected {AGENT_PROTOCOL_VERSION}"
2703 ),
2704 },
2705 )
2706 .await?;
2707 return Ok(());
2708 }
2709 AgentRequest::Hello { client_version, .. } => {
2710 send_response(
2711 &mut writer,
2712 &AgentResponse::Error {
2713 message: format!(
2714 "cache client {client_version} does not match agent {}",
2715 self.version
2716 ),
2717 },
2718 )
2719 .await?;
2720 return Ok(());
2721 }
2722 _ => bail!("the first agent request must be hello"),
2723 }
2724 send_response(
2725 &mut writer,
2726 &AgentResponse::Hello {
2727 protocol: AGENT_PROTOCOL_VERSION,
2728 agent_version: self.version.to_string(),
2729 },
2730 )
2731 .await?;
2732
2733 let mut connection = ConnectionUploads::default();
2737 while let Some(line) = read_request(&mut reader).await? {
2738 let response = match serde_json::from_str(&line) {
2739 Ok(request) => self.respond_on(request, &mut connection).await,
2740 Err(error) => AgentResponse::Error {
2741 message: format!("invalid agent request: {error}"),
2742 },
2743 };
2744 send_response(&mut writer, &response).await?;
2745 }
2746 Ok(())
2747 }
2748}
2749
2750async fn read_request<R>(reader: &mut R) -> Result<Option<String>>
2756where
2757 R: AsyncBufRead + Unpin,
2758{
2759 let mut line = Vec::new();
2760 loop {
2761 let available = reader.fill_buf().await?;
2762 if available.is_empty() {
2763 break;
2764 }
2765 let (consumed, complete) = match available.iter().position(|byte| *byte == b'\n') {
2766 Some(index) => (index, true),
2767 None => (available.len(), false),
2768 };
2769 if line.len() + consumed > MAX_REQUEST_BYTES {
2770 bail!("agent request exceeded {MAX_REQUEST_BYTES} bytes");
2771 }
2772 line.extend_from_slice(&available[..consumed]);
2773 reader.consume(consumed + usize::from(complete));
2775 if complete {
2776 return Ok(Some(String::from_utf8(line)?));
2777 }
2778 }
2779 if line.is_empty() {
2780 Ok(None)
2781 } else {
2782 Ok(Some(String::from_utf8(line)?))
2783 }
2784}
2785
2786pub fn is_task_identity(task: &str) -> bool {
2792 task.len() == 64
2793 && task
2794 .bytes()
2795 .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
2796}
2797
2798fn validate_task_identity(task: &str) -> Result<()> {
2799 if !is_task_identity(task) {
2800 bail!("invalid task action identity");
2801 }
2802 Ok(())
2803}
2804
2805fn task_manifest_dir(store: &Path) -> PathBuf {
2807 store.join("task-manifests").join("v1")
2808}
2809
2810pub fn task_manifest_actions(store: &Path, task: &str) -> Result<Vec<CacheDigest>> {
2818 validate_task_identity(task)?;
2819 let path = task_manifest_dir(store).join(format!("{task}.json"));
2820 let bytes = match fs::read(&path) {
2821 Ok(bytes) => bytes,
2822 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
2823 Err(error) => {
2824 return Err(error).wrap_err_with(|| format!("failed to read {}", path.display()));
2825 }
2826 };
2827 let Ok(manifest) = serde_json::from_slice::<TaskActionManifest>(&bytes) else {
2828 return Ok(Vec::new());
2829 };
2830 if validate_task_manifest(&manifest, task).is_err() {
2831 return Ok(Vec::new());
2832 }
2833 Ok(manifest
2834 .predictions
2835 .into_iter()
2836 .map(|prediction| prediction.action)
2837 .collect())
2838}
2839
2840fn validate_action_prediction(prediction: &ActionPrediction) -> Result<()> {
2841 if prediction.validate() {
2842 Ok(())
2843 } else {
2844 bail!("invalid action prediction")
2845 }
2846}
2847
2848fn validate_task_manifest(manifest: &TaskActionManifest, task: &str) -> Result<()> {
2849 if manifest.task == task && manifest.validate() {
2850 Ok(())
2851 } else {
2852 bail!("invalid task action manifest")
2853 }
2854}
2855
2856fn merge_task_manifests(
2857 task: &str,
2858 base: Option<TaskActionManifest>,
2859 update: TaskActionManifest,
2860) -> Result<TaskActionManifest> {
2861 validate_task_manifest(&update, task)?;
2862 let mut predictions = BTreeMap::new();
2863 if let Some(base) = base {
2864 validate_task_manifest(&base, task)?;
2865 predictions.extend(
2866 base.predictions
2867 .into_iter()
2868 .map(|prediction| (prediction.invocation.clone(), prediction)),
2869 );
2870 }
2871 predictions.extend(
2872 update
2873 .predictions
2874 .into_iter()
2875 .map(|prediction| (prediction.invocation.clone(), prediction)),
2876 );
2877 let manifest = TaskActionManifest {
2878 version: TASK_ACTION_MANIFEST_VERSION,
2879 task: task.to_owned(),
2880 predictions: predictions.into_values().collect(),
2881 };
2882 validate_task_manifest(&manifest, task)?;
2883 Ok(manifest)
2884}
2885
2886fn merge_remote_task_manifest(
2887 task: &str,
2888 remote: TaskActionManifest,
2889 local: TaskActionManifest,
2890) -> (TaskActionManifest, bool) {
2891 match merge_task_manifests(task, Some(remote), local.clone()) {
2892 Ok(manifest) => (manifest, true),
2893 Err(error) => {
2894 warn!("remote task action manifest merge failed for {task}: {error}");
2895 (local, false)
2896 }
2897 }
2898}
2899
2900async fn send_response(
2901 writer: &mut (impl AsyncWrite + Unpin),
2902 response: &AgentResponse,
2903) -> Result<()> {
2904 let mut encoded = serde_json::to_vec(response)?;
2905 encoded.push(b'\n');
2906 writer.write_all(&encoded).await?;
2907 writer.flush().await?;
2908 Ok(())
2909}
2910
2911#[cfg(test)]
2912#[path = "agent_tests.rs"]
2913mod tests;