1use crate::{
2 ActionPrediction, BlobPackLimits, BlobSource, BlobUpload, CacheDigest, CacheDirectory,
3 LocalActionCache, LocalCas, MAX_STAGED_BLOB_PACK_BYTES, MAX_STAGED_BLOB_PACK_ITEMS,
4 ManifestPutOutcome, RemoteActionResult, RemoteCacheClient, RemoteCacheMode, RustcMetadata,
5 TaskActionManifest, blob_pack_chunk, canonical_json,
6};
7use eyre::{Context, Result, bail};
8use futures_util::{FutureExt, StreamExt, future::BoxFuture, stream};
9use log::warn;
10use serde::{Deserialize, Serialize};
11use std::collections::BTreeMap;
12use std::fs;
13use std::path::{Path, PathBuf};
14use std::sync::atomic::{AtomicU64, Ordering};
15use std::sync::{Arc, Mutex, Weak};
16use std::time::{Duration, Instant};
17use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
18
19const MAX_EXECUTABLE_IDENTITIES: usize = 64;
20const MAX_EXECUTABLE_IDENTITY_SIZE: usize = 64 * 1024;
21const MAX_EXECUTABLE_IDENTITY_BYTES: usize = 256 * 1024;
22const TASK_ACTION_MANIFEST_VERSION: u8 = 1;
23const MAX_TASK_ACTION_PREDICTIONS: usize = 16 * 1024;
24const MAX_REMOTE_TRANSFERS: usize = 64;
25const MAX_PREFETCH_TRANSFERS: usize = 48;
26const MAX_PREFETCH_ACTION_BATCH: usize = 256;
27const PREFETCH_ACTION_BATCH_DELAY: Duration = Duration::from_millis(5);
28const MAX_PREFETCH_DIRECTORY_OBJECTS: usize = 100_000;
29const MAX_PREFETCH_OBJECTS_PER_WAVE: usize = 100_000;
30const DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES: u64 = 5 * 1024 * 1024 * 1024;
31
32pub struct AgentRemoteCache {
34 pub client: RemoteCacheClient,
36 pub mode: RemoteCacheMode,
38 pub staging_dir: PathBuf,
40}
41
42pub const AGENT_PROTOCOL_VERSION: u8 = 2;
44const MAX_REQUEST_BYTES: usize = 16 * 1024 * 1024;
49
50#[derive(Debug, Clone, Serialize, Deserialize)]
52#[serde(tag = "type", rename_all = "snake_case")]
53pub enum AgentRequest {
54 Hello {
56 protocol: u8,
58 client_version: String,
60 },
61 FindBlob {
63 digest: CacheDigest,
65 },
66 FindBlobs {
68 digests: Vec<CacheDigest>,
70 },
71 StoreBlob {
73 digest: CacheDigest,
75 source: PathBuf,
77 },
78 FindActionResult {
80 action: CacheDigest,
82 },
83 RecordActionHit {
85 action: CacheDigest,
87 restore: RestoreStats,
89 },
90 RecordBypass {
92 kind: String,
94 },
95 RecordUnconsulted,
98 RecordCompilerInvocation {
100 outcome: String,
102 crate_name: Option<String>,
104 duration_ns: u64,
106 },
107 RecordActionVerification {
109 matched: bool,
111 restore: RestoreStats,
113 },
114 StoreActionResult {
116 result: RemoteActionResult,
118 },
119 FindActionPrediction {
121 task: String,
123 invocation: CacheDigest,
125 },
126 RecordActionPrediction {
128 task: String,
130 prediction: ActionPrediction,
132 },
133 FindExecutableIdentity {
135 executable: PathBuf,
137 environment: BTreeMap<String, Option<String>>,
139 },
140 StoreExecutableIdentity {
142 executable: PathBuf,
144 environment: BTreeMap<String, Option<String>>,
146 stdout: Vec<u8>,
148 },
149}
150
151#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
153#[serde(deny_unknown_fields)]
154pub struct RestoreStats {
155 pub duration_ns: u64,
157 pub avoided_compiler_duration_ns: u64,
160 pub output_files: u64,
162 pub output_bytes: u64,
164 pub reflinked_output_files: u64,
166 pub reflinked_output_bytes: u64,
168 pub copied_output_files: u64,
170 pub copied_output_bytes: u64,
172}
173
174#[derive(Debug, Clone, Serialize, Deserialize)]
176#[serde(tag = "type", rename_all = "snake_case")]
177pub enum AgentResponse {
178 Hello {
180 protocol: u8,
182 agent_version: String,
184 },
185 Blob {
187 path: Option<PathBuf>,
189 },
190 Blobs {
192 paths: Vec<Option<PathBuf>>,
194 },
195 Stored {
197 path: PathBuf,
199 },
200 ActionResult {
202 result: Option<RemoteActionResult>,
204 },
205 ActionHitRecorded,
207 ActionVerificationRecorded,
209 BypassRecorded,
211 UnconsultedRecorded,
213 CompilerInvocationRecorded,
215 ActionStored {
217 path: PathBuf,
219 },
220 ActionPrediction {
222 prediction: Option<ActionPrediction>,
224 },
225 ActionPredictionRecorded,
227 ExecutableIdentity {
229 stdout: Option<Vec<u8>>,
231 },
232 Error {
234 message: String,
236 },
237}
238
239#[derive(Debug, Clone, Default, PartialEq, Eq)]
246#[non_exhaustive]
247pub struct AgentStats {
248 pub session_duration_ns: u64,
250 pub lookups: u64,
252 pub unconsulted: u64,
258 pub hits: u64,
260 pub stores: u64,
262 pub stored_bytes: u64,
264 pub verifications: u64,
266 pub divergences: u64,
268 pub downloaded_bytes: u64,
270 pub uploaded_bytes: u64,
272 pub prefetched_actions: u64,
274 pub bypasses: BTreeMap<String, u64>,
276 pub avoided_compiler_duration_ns: u64,
278 pub compiler: BTreeMap<String, CompilerStats>,
280 pub slow_compilations: BTreeMap<String, u64>,
282 pub remote_failures: u64,
290 pub remote_manifest_lookups: u64,
292 pub remote_manifest_lookup_duration_ns: u64,
294 pub remote_action_lookups: u64,
296 pub remote_action_lookup_duration_ns: u64,
298 pub remote_blob_requests: u64,
300 pub remote_blob_pack_requests: u64,
302 pub remote_blob_pack_blobs: u64,
304 pub remote_blob_transfer_duration_ns: u64,
306 pub local_cas_write_duration_ns: u64,
308 pub prefetch_runs: u64,
310 pub prefetch_duration_ns: u64,
312 pub materialization_duration_ns: u64,
314 pub restored_output_files: u64,
316 pub restored_output_bytes: u64,
318 pub reflinked_output_files: u64,
320 pub reflinked_output_bytes: u64,
322 pub copied_output_files: u64,
324 pub copied_output_bytes: u64,
326}
327
328#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
334#[non_exhaustive]
335pub struct CompilerStats {
336 pub invocations: u64,
338 pub duration_ns: u64,
340}
341
342impl CompilerStats {
343 pub fn new(invocations: u64, duration_ns: u64) -> Self {
345 Self {
346 invocations,
347 duration_ns,
348 }
349 }
350}
351
352#[derive(Default)]
353struct AtomicAgentStats {
354 lookups: AtomicU64,
355 unconsulted: AtomicU64,
356 hits: AtomicU64,
357 stores: AtomicU64,
358 stored_bytes: AtomicU64,
359 verifications: AtomicU64,
360 divergences: AtomicU64,
361 downloaded_bytes: AtomicU64,
362 uploaded_bytes: AtomicU64,
363 prefetched_actions: AtomicU64,
364 remote_failures: AtomicU64,
365 remote_manifest_lookups: AtomicU64,
366 remote_manifest_lookup_duration_ns: AtomicU64,
367 remote_action_lookups: AtomicU64,
368 remote_action_lookup_duration_ns: AtomicU64,
369 remote_blob_requests: AtomicU64,
370 remote_blob_pack_requests: AtomicU64,
371 remote_blob_pack_blobs: AtomicU64,
372 remote_blob_transfer_duration_ns: AtomicU64,
373 local_cas_write_duration_ns: AtomicU64,
374 prefetch_runs: AtomicU64,
375 prefetch_duration_ns: AtomicU64,
376 materialization_duration_ns: AtomicU64,
377 bypasses: Mutex<BTreeMap<String, u64>>,
378 avoided_compiler_duration_ns: AtomicU64,
379 compiler: Mutex<BTreeMap<String, CompilerStats>>,
380 slow_compilations: Mutex<BTreeMap<String, u64>>,
381 restored_output_files: AtomicU64,
382 restored_output_bytes: AtomicU64,
383 reflinked_output_files: AtomicU64,
384 reflinked_output_bytes: AtomicU64,
385 copied_output_files: AtomicU64,
386 copied_output_bytes: AtomicU64,
387}
388
389struct AtomicDurationTimer<'a> {
390 started: Instant,
391 target: &'a AtomicU64,
392}
393
394impl<'a> AtomicDurationTimer<'a> {
395 fn start(target: &'a AtomicU64) -> Self {
396 Self {
397 started: Instant::now(),
398 target,
399 }
400 }
401}
402
403impl Drop for AtomicDurationTimer<'_> {
404 fn drop(&mut self) {
405 atomic_saturating_add(self.target, duration_ns(self.started));
406 }
407}
408
409fn duration_ns(started: Instant) -> u64 {
410 started.elapsed().as_nanos().try_into().unwrap_or(u64::MAX)
411}
412
413fn atomic_saturating_add(target: &AtomicU64, value: u64) {
414 let _ = target.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
415 Some(current.saturating_add(value))
416 });
417}
418
419fn queue_prefetch_digest(
420 verified: &BTreeMap<CacheDigest, PathBuf>,
421 pending: &mut BTreeMap<CacheDigest, ()>,
422 digest: CacheDigest,
423) {
424 if verified.contains_key(&digest) || pending.contains_key(&digest) {
425 return;
426 }
427 pending.insert(digest, ());
428}
429
430fn queue_prefetch_directory(
431 seen: &BTreeMap<CacheDigest, ()>,
432 pending: &mut BTreeMap<CacheDigest, ()>,
433 digest: CacheDigest,
434 limit: usize,
435) -> bool {
436 if seen.contains_key(&digest) || pending.contains_key(&digest) {
437 return true;
438 }
439 if seen.len().saturating_add(pending.len()) >= limit {
440 return false;
441 }
442 pending.insert(digest, ());
443 true
444}
445
446#[derive(Clone)]
452pub struct CacheAgent {
453 cas: LocalCas,
454 actions: LocalActionCache,
455 verified_blobs: Arc<Mutex<BTreeMap<CacheDigest, PathBuf>>>,
456 version: Arc<str>,
457 write_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
458 action_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
459 stats: Arc<AtomicAgentStats>,
460 executable_identities: Arc<Mutex<BTreeMap<ExecutableIdentityKey, Vec<u8>>>>,
461 manifest_dir: Arc<PathBuf>,
462 task_actions: Arc<Mutex<BTreeMap<String, TaskActionState>>>,
463 next_task_run: Arc<AtomicU64>,
464 manifest_write_lock: Arc<Mutex<()>>,
465 remote: Option<Arc<RemoteCacheClient>>,
466 remote_mode: RemoteCacheMode,
467 remote_staging_dir: Arc<PathBuf>,
468 remote_download_limit: u64,
469 remote_download_bytes: Arc<AtomicU64>,
470 pending_remote_actions: Arc<Mutex<BTreeMap<CacheDigest, RemoteActionResult>>>,
471 remote_transfers: Arc<tokio::sync::Semaphore>,
472 prefetch_transfers: Arc<tokio::sync::Semaphore>,
473 prefetch_tasks: Arc<Mutex<Vec<tokio::task::JoinHandle<()>>>>,
474}
475
476#[derive(Debug, Clone, Default)]
477struct TaskActionState {
478 manifest: String,
479 baseline_loaded: bool,
480 predictions: BTreeMap<CacheDigest, ActionPrediction>,
481 pending_predictions: BTreeMap<CacheDigest, ActionPrediction>,
482 remote_etag: Option<String>,
483}
484
485struct PrefetchedAction {
486 adapter: String,
487 result: RemoteActionResult,
488}
489
490struct RemoteDownloadReservation {
491 counter: Arc<AtomicU64>,
492 reserved: u64,
493 committed: bool,
494}
495
496impl RemoteDownloadReservation {
497 fn bytes(&self) -> u64 {
498 self.reserved
499 }
500
501 fn commit(mut self, bytes: u64) {
502 debug_assert!(bytes <= self.reserved);
503 self.counter
504 .fetch_sub(self.reserved.saturating_sub(bytes), Ordering::AcqRel);
505 self.committed = true;
506 }
507}
508
509impl Drop for RemoteDownloadReservation {
510 fn drop(&mut self) {
511 if !self.committed {
512 self.counter.fetch_sub(self.reserved, Ordering::AcqRel);
513 }
514 }
515}
516
517#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
518struct ExecutableIdentityKey {
519 executable: PathBuf,
520 environment: BTreeMap<String, Option<String>>,
521}
522
523impl CacheAgent {
524 pub fn new(cache_dir: impl Into<PathBuf>, version: impl Into<Arc<str>>) -> Self {
526 Self::build(cache_dir.into(), version.into(), None, 0)
527 }
528
529 pub fn new_remote(
531 cache_dir: impl Into<PathBuf>,
532 version: impl Into<Arc<str>>,
533 remote: AgentRemoteCache,
534 ) -> Self {
535 Self::build(
536 cache_dir.into(),
537 version.into(),
538 Some(remote),
539 DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES,
540 )
541 }
542
543 pub fn new_remote_with_download_limit(
545 cache_dir: impl Into<PathBuf>,
546 version: impl Into<Arc<str>>,
547 remote: AgentRemoteCache,
548 max_remote_download_bytes: u64,
549 ) -> Self {
550 Self::build(
551 cache_dir.into(),
552 version.into(),
553 Some(remote),
554 max_remote_download_bytes,
555 )
556 }
557
558 fn build(
559 cache_dir: PathBuf,
560 version: Arc<str>,
561 remote: Option<AgentRemoteCache>,
562 remote_download_limit: u64,
563 ) -> Self {
564 let remote_mode = remote
565 .as_ref()
566 .map_or(RemoteCacheMode::ReadOnly, |remote| remote.mode);
567 let remote_staging_dir = remote.as_ref().map_or_else(
568 || cache_dir.join("remote"),
569 |remote| remote.staging_dir.clone(),
570 );
571 let remote = remote.map(|remote| Arc::new(remote.client));
572 Self {
573 cas: LocalCas::new(cache_dir.clone()),
574 actions: LocalActionCache::new(cache_dir.clone()),
575 verified_blobs: Arc::new(Mutex::new(BTreeMap::new())),
576 version,
577 write_locks: Arc::new(Mutex::new(BTreeMap::new())),
578 action_locks: Arc::new(Mutex::new(BTreeMap::new())),
579 stats: Arc::new(AtomicAgentStats::default()),
580 executable_identities: Arc::new(Mutex::new(BTreeMap::new())),
581 manifest_dir: Arc::new(task_manifest_dir(&cache_dir)),
582 task_actions: Arc::new(Mutex::new(BTreeMap::new())),
583 next_task_run: Arc::new(AtomicU64::new(0)),
584 manifest_write_lock: Arc::new(Mutex::new(())),
585 remote,
586 remote_mode,
587 remote_staging_dir: Arc::new(remote_staging_dir),
588 remote_download_limit,
589 remote_download_bytes: Arc::new(AtomicU64::new(0)),
590 pending_remote_actions: Arc::new(Mutex::new(BTreeMap::new())),
591 remote_transfers: Arc::new(tokio::sync::Semaphore::new(MAX_REMOTE_TRANSFERS)),
592 prefetch_transfers: Arc::new(tokio::sync::Semaphore::new(MAX_PREFETCH_TRANSFERS)),
593 prefetch_tasks: Arc::new(Mutex::new(Vec::new())),
594 }
595 }
596
597 fn reserve_remote_download(&self, bytes: u64) -> Result<RemoteDownloadReservation> {
598 let mut current = self.remote_download_bytes.load(Ordering::Acquire);
599 loop {
600 let next = current
601 .checked_add(bytes)
602 .ok_or_else(|| eyre::eyre!("remote cache download budget overflowed"))?;
603 if next > self.remote_download_limit {
604 bail!(
605 "remote cache download budget exceeded: {} bytes requested with {} of {} bytes already used",
606 bytes,
607 current,
608 self.remote_download_limit
609 );
610 }
611 match self.remote_download_bytes.compare_exchange_weak(
612 current,
613 next,
614 Ordering::AcqRel,
615 Ordering::Acquire,
616 ) {
617 Ok(_) => {
618 return Ok(RemoteDownloadReservation {
619 counter: self.remote_download_bytes.clone(),
620 reserved: bytes,
621 committed: false,
622 });
623 }
624 Err(observed) => current = observed,
625 }
626 }
627 }
628
629 fn reserve_remote_download_up_to(&self, requested: u64) -> Result<RemoteDownloadReservation> {
630 let mut current = self.remote_download_bytes.load(Ordering::Acquire);
631 loop {
632 let reserved = requested.min(self.remote_download_limit.saturating_sub(current));
633 let next = current + reserved;
634 match self.remote_download_bytes.compare_exchange_weak(
635 current,
636 next,
637 Ordering::AcqRel,
638 Ordering::Acquire,
639 ) {
640 Ok(_) => {
641 return Ok(RemoteDownloadReservation {
642 counter: self.remote_download_bytes.clone(),
643 reserved,
644 committed: false,
645 });
646 }
647 Err(observed) => current = observed,
648 }
649 }
650 }
651
652 pub async fn begin_task(&self, task: &str) -> Result<String> {
654 self.begin_task_with_remote_errors(task, false).await
655 }
656
657 pub async fn prefetch_task(&self, task: &str) -> Result<String> {
659 let run = self.begin_task_with_remote_errors(task, true).await?;
660 self.wait_for_prefetches().await;
661 Ok(run)
662 }
663
664 async fn begin_task_with_remote_errors(&self, task: &str, strict: bool) -> Result<String> {
665 validate_task_identity(task)?;
666 let (remote_manifest, mut remote_etag) = if self.remote_mode.reads() {
667 match self.get_remote_task_manifest(task).await {
668 Ok(Some((manifest, etag))) => (Some(manifest), Some(etag)),
669 Ok(None) => (None, None),
670 Err(error) => {
671 if strict {
672 return Err(error).wrap_err_with(|| {
673 format!("remote task action manifest lookup failed for {task}")
674 });
675 }
676 self.note_remote_failure();
677 warn!("remote task action manifest lookup failed for {task}: {error}");
678 (None, None)
679 }
680 }
681 } else {
682 (None, None)
683 };
684 let manifest = {
685 let _write_guard = self.manifest_write_lock.lock().unwrap();
686 let _file_guard = self.lock_task_manifest(task)?;
687 let local_manifest = self.load_task_manifest(task)?;
688 let manifest = match (remote_manifest, local_manifest) {
689 (Some(remote), Some(local)) => {
690 let (manifest, merged) = merge_remote_task_manifest(task, remote, local);
691 if !merged {
692 remote_etag = None;
693 }
694 Some(manifest)
695 }
696 (Some(remote), None) => Some(remote),
697 (None, local) => local,
698 };
699 if let Some(manifest) = &manifest {
700 self.persist_task_manifest(manifest)?;
701 }
702 manifest
703 };
704 let state = if let Some(manifest) = manifest {
705 TaskActionState {
706 manifest: task.to_string(),
707 baseline_loaded: true,
708 predictions: manifest
709 .predictions
710 .into_iter()
711 .map(|prediction| (prediction.invocation.clone(), prediction))
712 .collect(),
713 pending_predictions: BTreeMap::new(),
714 remote_etag,
715 }
716 } else {
717 TaskActionState {
718 manifest: task.to_string(),
719 baseline_loaded: true,
720 remote_etag,
721 ..TaskActionState::default()
722 }
723 };
724 let sequence = self.next_task_run.fetch_add(1, Ordering::Relaxed);
725 let run =
726 CacheDigest::blake3(format!("{task}\0{}\0{sequence}", std::process::id()).as_bytes())
727 .hash;
728 let predictions = state.predictions.values().cloned().collect();
729 self.task_actions.lock().unwrap().insert(run.clone(), state);
730 self.spawn_prefetch_predictions(predictions);
731 Ok(run)
732 }
733
734 pub async fn cancel_prefetches(&self) {
736 let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
737 for task in &tasks {
738 task.abort();
739 }
740 for task in tasks {
741 if let Err(error) = task.await
742 && !error.is_cancelled()
743 {
744 warn!("remote action prefetch task failed: {error}");
745 }
746 }
747 }
748
749 async fn wait_for_prefetches(&self) {
750 let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
751 for task in tasks {
752 if let Err(error) = task.await {
753 warn!("remote action prefetch task failed: {error}");
754 }
755 }
756 }
757
758 pub async fn commit_task(&self, run: &str) -> Result<()> {
760 validate_task_identity(run)?;
761 let state = self
762 .task_actions
763 .lock()
764 .unwrap()
765 .get(run)
766 .cloned()
767 .ok_or_else(|| eyre::eyre!("task action manifest baseline was not loaded"))?;
768 if !state.baseline_loaded {
769 bail!("task action manifest baseline was not loaded");
770 }
771 let task = state.manifest;
772 validate_task_identity(&task)?;
773 let manifest = {
774 let _write_guard = self.manifest_write_lock.lock().unwrap();
775 let _file_guard = self.lock_task_manifest(&task)?;
776 let mut predictions = self
777 .load_task_manifest(&task)?
778 .map(|manifest| {
779 manifest
780 .predictions
781 .into_iter()
782 .map(|prediction| (prediction.invocation.clone(), prediction))
783 .collect::<BTreeMap<_, _>>()
784 })
785 .unwrap_or_default();
786 predictions.extend(state.pending_predictions);
791 let manifest = TaskActionManifest {
792 version: TASK_ACTION_MANIFEST_VERSION,
793 task: task.clone(),
794 predictions: predictions.into_values().collect(),
795 };
796 validate_task_manifest(&manifest, &task)?;
797 self.persist_task_manifest(&manifest)?;
798 manifest
799 };
800 self.task_actions.lock().unwrap().remove(run);
801 if self.remote_mode.writes() {
802 match self
803 .put_remote_task_manifest(&task, manifest, state.remote_etag)
804 .await
805 {
806 Ok(remote_manifest) => {
807 let _write_guard = self.manifest_write_lock.lock().unwrap();
808 let reconciliation = (|| {
809 let _file_guard = self.lock_task_manifest(&task)?;
810 let manifest = match self.load_task_manifest(&task)? {
811 Some(local) => {
812 merge_remote_task_manifest(&task, remote_manifest, local).0
813 }
814 None => remote_manifest,
815 };
816 self.persist_task_manifest(&manifest)
817 })();
818 if let Err(error) = reconciliation {
819 warn!(
820 "remote task action manifest reconciliation failed for {task}: {error}"
821 );
822 }
823 }
824 Err(error) => {
825 self.note_remote_failure();
826 warn!("remote task action manifest upload failed for {task}: {error}");
827 }
828 }
829 }
830 Ok(())
831 }
832
833 fn task_manifest_path(&self, task: &str) -> PathBuf {
834 self.manifest_dir.join(format!("{task}.json"))
835 }
836
837 fn task_manifest_lock_path(&self, task: &str) -> PathBuf {
838 self.manifest_dir.join("locks").join(format!("{task}.lock"))
839 }
840
841 fn lock_task_manifest(&self, task: &str) -> Result<fslock::LockFile> {
842 let path = self.task_manifest_lock_path(task);
843 fs::create_dir_all(path.parent().expect("task manifest lock has a parent"))?;
844 let mut lock = fslock::LockFile::open(&path)?;
845 lock.lock()?;
846 Ok(lock)
847 }
848
849 fn load_task_manifest(&self, task: &str) -> Result<Option<TaskActionManifest>> {
850 match fs::read(self.task_manifest_path(task)) {
851 Ok(contents) => Ok(Some(self.parse_task_manifest(task, &contents, false)?)),
852 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
853 Err(error) => Err(error.into()),
854 }
855 }
856
857 fn parse_task_manifest(
858 &self,
859 task: &str,
860 contents: &[u8],
861 require_canonical: bool,
862 ) -> Result<TaskActionManifest> {
863 let manifest: TaskActionManifest = serde_json::from_slice(contents)?;
864 validate_task_manifest(&manifest, task)?;
865 if require_canonical && canonical_json(&manifest)? != contents {
866 bail!("task action manifest is not canonical JSON");
867 }
868 Ok(manifest)
869 }
870
871 fn task_manifest_selector(task: &str) -> Result<(Vec<u8>, CacheDigest)> {
872 TaskActionManifest::selector(task)
873 }
874
875 fn persist_task_manifest(&self, manifest: &TaskActionManifest) -> Result<()> {
876 let bytes = canonical_json(manifest)?;
877 fs::create_dir_all(self.manifest_dir.as_path())?;
878 let mut temporary = tempfile::NamedTempFile::new_in(self.manifest_dir.as_path())?;
879 std::io::Write::write_all(temporary.as_file_mut(), &bytes)?;
880 temporary.as_file_mut().sync_all()?;
881 temporary
882 .persist(self.task_manifest_path(&manifest.task))
883 .map_err(|error| error.error)?;
884 Ok(())
885 }
886
887 async fn get_remote_task_manifest(
888 &self,
889 task: &str,
890 ) -> Result<Option<(TaskActionManifest, String)>> {
891 let Some(remote) = &self.remote else {
892 return Ok(None);
893 };
894 let (_, selector) = Self::task_manifest_selector(task)?;
895 let _permit = self.remote_transfers.acquire().await?;
896 self.stats
897 .remote_manifest_lookups
898 .fetch_add(1, Ordering::Relaxed);
899 let _timer = AtomicDurationTimer::start(&self.stats.remote_manifest_lookup_duration_ns);
900 let Some(remote_manifest) = remote.get_action_manifest(&selector).await? else {
901 return Ok(None);
902 };
903 let manifest = self.parse_task_manifest(task, &remote_manifest.bytes, true)?;
904 Ok(Some((manifest, remote_manifest.etag)))
905 }
906
907 async fn put_remote_task_manifest(
908 &self,
909 task: &str,
910 mut manifest: TaskActionManifest,
911 mut expected_etag: Option<String>,
912 ) -> Result<TaskActionManifest> {
913 let Some(remote) = &self.remote else {
914 return Ok(manifest);
915 };
916 let (_, selector) = Self::task_manifest_selector(task)?;
917 for _ in 0..4 {
918 let bytes = canonical_json(&manifest)?;
919 let outcome = {
920 let _permit = self.remote_transfers.acquire().await?;
921 remote
922 .put_action_manifest(&selector, &bytes, expected_etag.as_deref())
923 .await?
924 };
925 match outcome {
926 ManifestPutOutcome::Stored => return Ok(manifest),
927 ManifestPutOutcome::PreconditionFailed => {
928 let Some((remote_manifest, etag)) = self.get_remote_task_manifest(task).await?
929 else {
930 expected_etag = None;
931 continue;
932 };
933 manifest = merge_task_manifests(task, Some(remote_manifest), manifest)?;
934 expected_etag = Some(etag);
935 }
936 }
937 }
938 bail!("remote task action manifest changed too frequently")
939 }
940
941 pub fn stats(&self) -> AgentStats {
943 AgentStats {
944 session_duration_ns: 0,
945 lookups: self.stats.lookups.load(Ordering::Relaxed),
946 unconsulted: self.stats.unconsulted.load(Ordering::Relaxed),
947 hits: self.stats.hits.load(Ordering::Relaxed),
948 stores: self.stats.stores.load(Ordering::Relaxed),
949 stored_bytes: self.stats.stored_bytes.load(Ordering::Relaxed),
950 verifications: self.stats.verifications.load(Ordering::Relaxed),
951 divergences: self.stats.divergences.load(Ordering::Relaxed),
952 downloaded_bytes: self.stats.downloaded_bytes.load(Ordering::Relaxed),
953 uploaded_bytes: self.stats.uploaded_bytes.load(Ordering::Relaxed),
954 prefetched_actions: self.stats.prefetched_actions.load(Ordering::Relaxed),
955 bypasses: self.stats.bypasses.lock().unwrap().clone(),
956 avoided_compiler_duration_ns: self
957 .stats
958 .avoided_compiler_duration_ns
959 .load(Ordering::Relaxed),
960 compiler: self.stats.compiler.lock().unwrap().clone(),
961 slow_compilations: self.stats.slow_compilations.lock().unwrap().clone(),
962 remote_failures: self.stats.remote_failures.load(Ordering::Relaxed),
963 remote_manifest_lookups: self.stats.remote_manifest_lookups.load(Ordering::Relaxed),
964 remote_manifest_lookup_duration_ns: self
965 .stats
966 .remote_manifest_lookup_duration_ns
967 .load(Ordering::Relaxed),
968 remote_action_lookups: self.stats.remote_action_lookups.load(Ordering::Relaxed),
969 remote_action_lookup_duration_ns: self
970 .stats
971 .remote_action_lookup_duration_ns
972 .load(Ordering::Relaxed),
973 remote_blob_requests: self.stats.remote_blob_requests.load(Ordering::Relaxed),
974 remote_blob_pack_requests: self.stats.remote_blob_pack_requests.load(Ordering::Relaxed),
975 remote_blob_pack_blobs: self.stats.remote_blob_pack_blobs.load(Ordering::Relaxed),
976 remote_blob_transfer_duration_ns: self
977 .stats
978 .remote_blob_transfer_duration_ns
979 .load(Ordering::Relaxed),
980 local_cas_write_duration_ns: self
981 .stats
982 .local_cas_write_duration_ns
983 .load(Ordering::Relaxed),
984 prefetch_runs: self.stats.prefetch_runs.load(Ordering::Relaxed),
985 prefetch_duration_ns: self.stats.prefetch_duration_ns.load(Ordering::Relaxed),
986 materialization_duration_ns: self
987 .stats
988 .materialization_duration_ns
989 .load(Ordering::Relaxed),
990 restored_output_files: self.stats.restored_output_files.load(Ordering::Relaxed),
991 restored_output_bytes: self.stats.restored_output_bytes.load(Ordering::Relaxed),
992 reflinked_output_files: self.stats.reflinked_output_files.load(Ordering::Relaxed),
993 reflinked_output_bytes: self.stats.reflinked_output_bytes.load(Ordering::Relaxed),
994 copied_output_files: self.stats.copied_output_files.load(Ordering::Relaxed),
995 copied_output_bytes: self.stats.copied_output_bytes.load(Ordering::Relaxed),
996 }
997 }
998
999 pub async fn handle_requests(
1006 &self,
1007 requests: impl IntoIterator<Item = AgentRequest>,
1008 ) -> Vec<AgentResponse> {
1009 let mut responses = Vec::new();
1010 for request in requests {
1011 responses.push(self.respond(request).await);
1012 }
1013 responses
1014 }
1015
1016 fn write_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1017 Self::digest_lock(&self.write_locks, digest)
1018 }
1019
1020 fn action_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1021 Self::digest_lock(&self.action_locks, digest)
1022 }
1023
1024 fn digest_lock(
1025 locks: &Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>,
1026 digest: &CacheDigest,
1027 ) -> Arc<tokio::sync::Mutex<()>> {
1028 let mut locks = locks.lock().unwrap();
1029 locks.retain(|_, lock| lock.strong_count() > 0);
1030 if let Some(lock) = locks.get(digest).and_then(Weak::upgrade) {
1031 return lock;
1032 }
1033 let lock = Arc::new(tokio::sync::Mutex::new(()));
1034 locks.insert(digest.clone(), Arc::downgrade(&lock));
1035 lock
1036 }
1037
1038 fn spawn_prefetch_predictions(&self, predictions: Vec<ActionPrediction>) {
1039 if predictions.is_empty() || !self.remote_mode.reads() || self.remote.is_none() {
1040 return;
1041 }
1042 let agent = self.clone();
1043 let task = tokio::spawn(async move {
1044 agent.prefetch_predictions(predictions.iter()).await;
1045 });
1046 self.prefetch_tasks.lock().unwrap().push(task);
1047 }
1048
1049 async fn prefetch_predictions<'a>(
1050 &self,
1051 predictions: impl Iterator<Item = &'a ActionPrediction>,
1052 ) {
1053 if !self.remote_mode.reads() || self.remote.is_none() {
1054 return;
1055 }
1056 self.stats.prefetch_runs.fetch_add(1, Ordering::Relaxed);
1057 let _timer = AtomicDurationTimer::start(&self.stats.prefetch_duration_ns);
1058 let mut actions = BTreeMap::new();
1059 for prediction in predictions {
1060 actions
1061 .entry(prediction.action.clone())
1062 .or_insert_with(|| prediction.adapter.clone());
1063 }
1064 let mut actions = actions.into_iter();
1065 let mut tasks = tokio::task::JoinSet::new();
1066 for _ in 0..MAX_PREFETCH_TRANSFERS {
1067 let Some((action, adapter)) = actions.next() else {
1068 break;
1069 };
1070 let agent = self.clone();
1071 tasks.spawn(async move { agent.resolve_prefetch_action(action, adapter).await });
1072 }
1073 let mut resolved = Vec::new();
1074 while !tasks.is_empty() {
1075 let result = if resolved.is_empty() {
1076 tasks.join_next().await
1077 } else {
1078 match tokio::time::timeout(PREFETCH_ACTION_BATCH_DELAY, tasks.join_next()).await {
1079 Ok(result) => result,
1080 Err(_) => {
1081 self.prefetch_resolved_actions(std::mem::take(&mut resolved))
1082 .await;
1083 continue;
1084 }
1085 }
1086 };
1087 let Some(result) = result else {
1088 break;
1089 };
1090 match result {
1091 Ok(Ok(Some(action))) => resolved.push(action),
1092 Ok(Ok(None)) => {}
1093 Ok(Err(error)) => {
1094 self.note_remote_failure();
1095 warn!("remote action prefetch failed: {error}");
1096 }
1097 Err(error) => {
1098 self.note_remote_failure();
1099 warn!("remote action prefetch task failed: {error}");
1100 }
1101 }
1102 if let Some((action, adapter)) = actions.next() {
1103 let agent = self.clone();
1104 tasks.spawn(async move { agent.resolve_prefetch_action(action, adapter).await });
1105 }
1106 if resolved.len() == MAX_PREFETCH_ACTION_BATCH {
1107 self.prefetch_resolved_actions(std::mem::take(&mut resolved))
1108 .await;
1109 }
1110 }
1111 if !resolved.is_empty() {
1112 self.prefetch_resolved_actions(resolved).await;
1113 }
1114 }
1115
1116 #[cfg(test)]
1117 async fn prefetch_action(&self, action: CacheDigest, adapter: String) -> Result<()> {
1118 if let Some(action) = self.resolve_prefetch_action(action, adapter).await? {
1119 self.prefetch_resolved_actions(vec![action]).await;
1120 }
1121 Ok(())
1122 }
1123
1124 async fn resolve_prefetch_action(
1125 &self,
1126 action: CacheDigest,
1127 adapter: String,
1128 ) -> Result<Option<PrefetchedAction>> {
1129 let remote = self
1130 .remote
1131 .as_ref()
1132 .ok_or_else(|| eyre::eyre!("remote cache is not configured"))?;
1133 let result = {
1134 let lock = self.action_lock(&action);
1135 let _guard = lock.lock().await;
1136 if self.actions.find(&action)?.is_some() {
1137 return Ok(None);
1138 }
1139 if let Some(result) = self
1140 .pending_remote_actions
1141 .lock()
1142 .unwrap()
1143 .get(&action)
1144 .cloned()
1145 {
1146 result
1147 } else {
1148 let _prefetch_permit = self.prefetch_transfers.acquire().await?;
1149 let result = {
1150 let _permit = self.remote_transfers.acquire().await?;
1151 self.get_remote_action_result(remote, &action).await?
1152 };
1153 let Some(result) = result else {
1154 return Ok(None);
1155 };
1156 self.pending_remote_actions
1157 .lock()
1158 .unwrap()
1159 .insert(action.clone(), result.clone());
1160 result
1161 }
1162 };
1163 Ok(Some(PrefetchedAction { adapter, result }))
1164 }
1165
1166 fn prefetch_resolved_actions(&self, actions: Vec<PrefetchedAction>) -> BoxFuture<'_, ()> {
1167 self.prefetch_resolved_actions_inner(actions).boxed()
1168 }
1169
1170 async fn prefetch_resolved_actions_inner(&self, actions: Vec<PrefetchedAction>) {
1171 let Some(remote) = self.remote.as_deref() else {
1172 return;
1173 };
1174 if actions.is_empty() {
1175 return;
1176 }
1177
1178 let mut top_level = BTreeMap::new();
1179 for action in &actions {
1180 for digest in [
1181 Some(&action.result.action),
1182 action.result.metadata.as_ref(),
1183 action.result.output_root.as_ref(),
1184 ]
1185 .into_iter()
1186 .flatten()
1187 {
1188 top_level.insert(digest.clone(), ());
1189 }
1190 }
1191 let mut verified = self
1192 .fetch_remote_blobs(
1193 remote,
1194 top_level.into_keys().collect(),
1195 Some(&self.prefetch_transfers),
1196 )
1197 .await;
1198
1199 let mut next = BTreeMap::new();
1200 let mut pending_directories = BTreeMap::new();
1201 let mut parsed_directories = BTreeMap::new();
1202 let mut rustc_metadata = BTreeMap::new();
1203 for action in &actions {
1204 if action.adapter == "rustc"
1205 && let Some(metadata_digest) = &action.result.metadata
1206 {
1207 match verified
1208 .get(metadata_digest)
1209 .ok_or_else(|| eyre::eyre!("remote rustc action metadata is missing"))
1210 .and_then(|path| Self::parse_rustc_metadata(path))
1211 {
1212 Ok(metadata) => {
1213 queue_prefetch_digest(&verified, &mut next, metadata.stdout.clone());
1214 queue_prefetch_digest(&verified, &mut next, metadata.stderr.clone());
1215 rustc_metadata.insert(metadata_digest.clone(), metadata);
1216 }
1217 Err(error) => warn!(
1218 "remote rustc action metadata prefetch failed for {}: {error}",
1219 action.result.action.hash
1220 ),
1221 }
1222 }
1223 if let Some(output_root) = &action.result.output_root {
1224 pending_directories.insert(output_root.clone(), ());
1225 }
1226 }
1227
1228 let mut seen_directories = BTreeMap::new();
1229 loop {
1230 let mut following = BTreeMap::new();
1231 let mut directory_limit_exceeded = false;
1232 for digest in pending_directories.into_keys() {
1233 following.remove(&digest);
1234 if seen_directories.insert(digest.clone(), ()).is_some() {
1235 continue;
1236 }
1237 if seen_directories.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
1238 warn!("remote action output tree is too large to prefetch");
1239 following.clear();
1240 break;
1241 }
1242 match verified
1243 .get(&digest)
1244 .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))
1245 .and_then(|path| Self::parse_cache_directory(path))
1246 {
1247 Ok(directory) => {
1248 for file in &directory.files {
1249 queue_prefetch_digest(&verified, &mut next, file.digest.clone());
1250 if next.len() >= MAX_PREFETCH_OBJECTS_PER_WAVE {
1251 self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1252 .await;
1253 }
1254 }
1255 for child in &directory.directories {
1256 if !queue_prefetch_directory(
1257 &seen_directories,
1258 &mut following,
1259 child.digest.clone(),
1260 MAX_PREFETCH_DIRECTORY_OBJECTS,
1261 ) {
1262 warn!("remote action output tree is too large to prefetch");
1263 directory_limit_exceeded = true;
1264 break;
1265 }
1266 queue_prefetch_digest(&verified, &mut next, child.digest.clone());
1267 if next.len() >= MAX_PREFETCH_OBJECTS_PER_WAVE {
1268 self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1269 .await;
1270 }
1271 }
1272 parsed_directories.insert(digest, directory);
1273 }
1274 Err(error) => warn!(
1275 "remote action output directory prefetch failed for {}: {error}",
1276 digest.hash
1277 ),
1278 }
1279 if directory_limit_exceeded {
1280 following.clear();
1281 break;
1282 }
1283 }
1284 self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1285 .await;
1286 if following.is_empty() {
1287 break;
1288 }
1289 pending_directories = following;
1290 }
1291
1292 for action in actions {
1293 match Self::validate_prefetched_action(
1294 &action,
1295 &verified,
1296 &rustc_metadata,
1297 &parsed_directories,
1298 ) {
1299 Ok(()) => {
1300 if let Err(error) = self.actions.store(&action.result) {
1301 warn!(
1302 "remote action prefetch could not publish {}: {error}",
1303 action.result.action.hash
1304 );
1305 continue;
1306 }
1307 self.pending_remote_actions
1308 .lock()
1309 .unwrap()
1310 .remove(&action.result.action);
1311 self.stats
1312 .prefetched_actions
1313 .fetch_add(1, Ordering::Relaxed);
1314 }
1315 Err(error) => warn!(
1316 "remote action prefetch was incomplete for {}: {error}",
1317 action.result.action.hash
1318 ),
1319 }
1320 }
1321 }
1322
1323 async fn flush_prefetch_digest_batch(
1324 &self,
1325 remote: &RemoteCacheClient,
1326 verified: &mut BTreeMap<CacheDigest, PathBuf>,
1327 pending: &mut BTreeMap<CacheDigest, ()>,
1328 ) {
1329 if pending.is_empty() {
1330 return;
1331 }
1332 let digests = std::mem::take(pending).into_keys().collect();
1333 verified.extend(
1334 self.fetch_remote_blobs(remote, digests, Some(&self.prefetch_transfers))
1335 .await,
1336 );
1337 }
1338
1339 async fn fetch_remote_blobs(
1340 &self,
1341 remote: &RemoteCacheClient,
1342 digests: Vec<CacheDigest>,
1343 prefetch_limit: Option<&tokio::sync::Semaphore>,
1344 ) -> BTreeMap<CacheDigest, PathBuf> {
1345 let mut verified = BTreeMap::new();
1346 let mut missing = BTreeMap::new();
1347 for digest in digests {
1348 match self.find_verified_blob(&digest) {
1349 Ok(Some(path)) => {
1350 verified.insert(digest, path);
1351 }
1352 Ok(None) => {
1353 missing.insert(digest, ());
1354 }
1355 Err(error) => warn!(
1356 "local cache blob lookup failed for {}: {error}",
1357 digest.hash
1358 ),
1359 }
1360 }
1361 if missing.is_empty() {
1362 return verified;
1363 }
1364
1365 let mut pack_candidates = missing.clone();
1366 while !pack_candidates.is_empty() {
1367 let candidates = match blob_pack_chunk(
1368 &pack_candidates.keys().cloned().collect::<Vec<_>>(),
1369 BlobPackLimits {
1370 max_items: MAX_STAGED_BLOB_PACK_ITEMS,
1371 max_bytes: MAX_STAGED_BLOB_PACK_BYTES,
1372 },
1373 ) {
1374 Ok(candidates) if !candidates.is_empty() => candidates,
1375 Ok(_) => break,
1376 Err(error) => {
1377 warn!("remote cache blob pack skipped: {error}");
1378 break;
1379 }
1380 };
1381 let mut pack_guards = BTreeMap::new();
1385 for digest in candidates {
1386 let guard = self.write_lock(&digest).lock_owned().await;
1387 match self.find_verified_blob(&digest) {
1388 Ok(Some(path)) => {
1389 pack_candidates.remove(&digest);
1390 missing.remove(&digest);
1391 verified.insert(digest, path);
1392 }
1393 Ok(None) => {
1394 pack_guards.insert(digest, guard);
1395 }
1396 Err(error) => {
1397 warn!(
1398 "local cache blob lookup failed for {}: {error}",
1399 digest.hash
1400 );
1401 pack_guards.insert(digest, guard);
1402 }
1403 }
1404 }
1405 let requested = pack_guards.keys().cloned().collect::<Vec<_>>();
1406 if requested.is_empty() {
1407 continue;
1408 }
1409 let requested_bytes = requested
1410 .iter()
1411 .fold(0_u64, |total, digest| total.saturating_add(digest.size));
1412 let pack_reservation = match self
1413 .reserve_remote_download_up_to(requested_bytes.min(MAX_STAGED_BLOB_PACK_BYTES))
1414 {
1415 Ok(reservation) if reservation.bytes() > 0 => reservation,
1416 Ok(_) => break,
1417 Err(error) => {
1418 warn!("remote cache blob pack skipped: {error}");
1419 break;
1420 }
1421 };
1422 let (pack, transfer_duration_ns) = {
1423 let _prefetch_permit = match prefetch_limit {
1424 Some(limit) => match limit.acquire().await {
1425 Ok(permit) => Some(permit),
1426 Err(error) => {
1427 warn!(
1428 "remote cache blob pack could not acquire prefetch limit: {error}"
1429 );
1430 break;
1431 }
1432 },
1433 None => None,
1434 };
1435 let _transfer_permit = match self.remote_transfers.acquire().await {
1436 Ok(permit) => permit,
1437 Err(error) => {
1438 warn!("remote cache blob pack could not acquire transfer limit: {error}");
1439 break;
1440 }
1441 };
1442 let transfer_started = Instant::now();
1443 let pack = remote
1444 .get_blob_pack_with_limit(
1445 &requested,
1446 self.remote_staging_dir.as_path(),
1447 pack_reservation.bytes(),
1448 )
1449 .await;
1450 (pack, duration_ns(transfer_started))
1451 };
1452 let pack = match pack {
1453 Ok(Some(pack)) => pack,
1454 Ok(None) => break,
1455 Err(error) => {
1456 atomic_saturating_add(
1457 &self.stats.remote_blob_transfer_duration_ns,
1458 transfer_duration_ns,
1459 );
1460 warn!(
1461 "remote cache blob pack failed; falling back to individual blobs: {error}"
1462 );
1463 break;
1464 }
1465 };
1466 pack_reservation.commit(pack.payload_bytes);
1467 atomic_saturating_add(
1468 &self.stats.remote_blob_transfer_duration_ns,
1469 transfer_duration_ns,
1470 );
1471 atomic_saturating_add(&self.stats.remote_blob_pack_requests, pack.requests);
1472 atomic_saturating_add(&self.stats.remote_blob_pack_blobs, pack.blob_count);
1473 atomic_saturating_add(&self.stats.downloaded_bytes, pack.payload_bytes);
1474 if pack.requested.is_empty() {
1475 for digest in &requested {
1479 pack_candidates.remove(digest);
1480 }
1481 continue;
1482 }
1483 for digest in &pack.requested {
1484 pack_candidates.remove(digest);
1485 }
1486 let mut ingests = stream::iter(pack.blobs.into_iter().map(|(digest, source)| {
1487 let digest_for_result = digest.clone();
1488 let guard = pack_guards
1489 .remove(&digest)
1490 .expect("requested packed blob has a write lock");
1491 async move {
1492 (
1493 digest_for_result,
1494 self.ingest_packed_blob(digest, source, guard).await,
1495 )
1496 }
1497 }))
1498 .buffer_unordered(MAX_PREFETCH_TRANSFERS);
1499 while let Some((digest, result)) = ingests.next().await {
1500 match result {
1501 Ok(path) => {
1502 missing.remove(&digest);
1503 verified.insert(digest, path);
1504 }
1505 Err(error) => warn!(
1506 "remote cache packed blob ingest failed for {}: {error}",
1507 digest.hash
1508 ),
1509 }
1510 }
1511 }
1512
1513 let mut transfers = stream::iter(missing.into_keys().map(|digest| {
1514 let digest_for_result = digest.clone();
1515 async move {
1516 (
1517 digest_for_result,
1518 self.fetch_remote_blob_with_limit(remote, &digest, prefetch_limit)
1519 .await,
1520 )
1521 }
1522 }))
1523 .buffer_unordered(MAX_PREFETCH_TRANSFERS);
1524 while let Some((digest, result)) = transfers.next().await {
1525 match result {
1526 Ok(path) => {
1527 verified.insert(digest, path);
1528 }
1529 Err(error) => warn!(
1530 "remote cache blob prefetch failed for {}: {error}",
1531 digest.hash
1532 ),
1533 }
1534 }
1535 verified
1536 }
1537
1538 async fn ingest_packed_blob(
1539 &self,
1540 digest: CacheDigest,
1541 source: PathBuf,
1542 _guard: tokio::sync::OwnedMutexGuard<()>,
1543 ) -> Result<PathBuf> {
1544 let digest_size = digest.size;
1545 let agent = self.clone();
1546 let (path, stored, cas_duration_ns) = tokio::task::spawn_blocking(move || {
1547 if let Some(path) = agent.find_verified_blob(&digest)? {
1548 return Ok::<_, eyre::Report>((path, false, 0));
1549 }
1550 let cas_started = Instant::now();
1551 let path = agent.cas.store_verified_file(&digest, &source)?;
1552 let cas_duration_ns = duration_ns(cas_started);
1553 agent.remember_verified_blob(&digest, &path);
1554 Ok((path, true, cas_duration_ns))
1555 })
1556 .await??;
1557 atomic_saturating_add(&self.stats.local_cas_write_duration_ns, cas_duration_ns);
1558 if stored {
1559 self.stats.stores.fetch_add(1, Ordering::Relaxed);
1560 atomic_saturating_add(&self.stats.stored_bytes, digest_size);
1561 }
1562 Ok(path)
1563 }
1564
1565 fn parse_rustc_metadata(path: &Path) -> Result<RustcMetadata> {
1566 let bytes = fs::read(path)?;
1567 let metadata: RustcMetadata = serde_json::from_slice(&bytes)?;
1568 if metadata.version != 1 || metadata.kind != "rustc" || canonical_json(&metadata)? != bytes
1569 {
1570 bail!("remote rustc action metadata is invalid");
1571 }
1572 Ok(metadata)
1573 }
1574
1575 fn parse_cache_directory(path: &Path) -> Result<CacheDirectory> {
1576 let bytes = fs::read(path)?;
1577 let directory: CacheDirectory = serde_json::from_slice(&bytes)?;
1578 if directory.version != 1 || canonical_json(&directory)? != bytes {
1579 bail!("remote action output directory is invalid");
1580 }
1581 Ok(directory)
1582 }
1583
1584 #[cfg(test)]
1585 fn load_cache_directory(&self, digest: &CacheDigest) -> Result<CacheDirectory> {
1586 let path = self
1587 .find_verified_blob(digest)?
1588 .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))?;
1589 Self::parse_cache_directory(&path)
1590 }
1591
1592 fn validate_prefetched_action(
1593 action: &PrefetchedAction,
1594 verified: &BTreeMap<CacheDigest, PathBuf>,
1595 rustc_metadata: &BTreeMap<CacheDigest, RustcMetadata>,
1596 directories: &BTreeMap<CacheDigest, CacheDirectory>,
1597 ) -> Result<()> {
1598 if !verified.contains_key(&action.result.action) {
1599 bail!("remote action descriptor is missing");
1600 }
1601 if let Some(metadata) = &action.result.metadata {
1602 if action.adapter == "rustc" {
1603 let metadata = rustc_metadata
1604 .get(metadata)
1605 .ok_or_else(|| eyre::eyre!("remote rustc action metadata is missing"))?;
1606 for digest in [&metadata.stdout, &metadata.stderr] {
1607 if !verified.contains_key(digest) {
1608 bail!("remote rustc action diagnostic blob is missing");
1609 }
1610 }
1611 } else if !verified.contains_key(metadata) {
1612 bail!("remote action metadata is missing");
1613 }
1614 }
1615 let mut pending = action
1616 .result
1617 .output_root
1618 .iter()
1619 .cloned()
1620 .collect::<Vec<_>>();
1621 let mut seen = BTreeMap::new();
1622 while let Some(digest) = pending.pop() {
1623 if seen.insert(digest.clone(), ()).is_some() {
1624 continue;
1625 }
1626 if seen.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
1627 bail!("remote action output tree is too large");
1628 }
1629 let directory = directories
1630 .get(&digest)
1631 .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))?;
1632 for file in &directory.files {
1633 if !verified.contains_key(&file.digest) {
1634 bail!("remote action output file is missing");
1635 }
1636 }
1637 pending.extend(
1638 directory
1639 .directories
1640 .iter()
1641 .map(|directory| directory.digest.clone()),
1642 );
1643 }
1644 Ok(())
1645 }
1646
1647 #[cfg(test)]
1648 async fn prefetch_output_tree(
1649 &self,
1650 remote: &RemoteCacheClient,
1651 output_root: &CacheDigest,
1652 ) -> Result<()> {
1653 let mut pending = vec![output_root.clone()];
1654 let mut seen = BTreeMap::new();
1655 while let Some(digest) = pending.pop() {
1656 if seen.insert(digest.clone(), ()).is_some() {
1657 continue;
1658 }
1659 if seen.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
1660 bail!("remote action output tree is too large");
1661 }
1662 self.fetch_remote_blob_with_limit(remote, &digest, Some(&self.prefetch_transfers))
1663 .await?;
1664 let directory = self.load_cache_directory(&digest)?;
1665 let mut transfers = stream::iter(directory.files.into_iter().map(|file| async move {
1666 self.fetch_remote_blob_with_limit(
1667 remote,
1668 &file.digest,
1669 Some(&self.prefetch_transfers),
1670 )
1671 .await
1672 .map(|_| ())
1673 }))
1674 .buffer_unordered(MAX_PREFETCH_TRANSFERS);
1675 while let Some(result) = transfers.next().await {
1676 result?;
1677 }
1678 pending.extend(
1679 directory
1680 .directories
1681 .into_iter()
1682 .map(|directory| directory.digest),
1683 );
1684 }
1685 Ok(())
1686 }
1687
1688 async fn fetch_remote_blob(
1689 &self,
1690 remote: &RemoteCacheClient,
1691 digest: &CacheDigest,
1692 ) -> Result<PathBuf> {
1693 self.fetch_remote_blob_with_limit(remote, digest, None)
1694 .await
1695 }
1696
1697 async fn fetch_remote_blob_with_limit(
1698 &self,
1699 remote: &RemoteCacheClient,
1700 digest: &CacheDigest,
1701 prefetch_limit: Option<&tokio::sync::Semaphore>,
1702 ) -> Result<PathBuf> {
1703 let lock = self.write_lock(digest);
1704 let _guard = lock.lock().await;
1705 if let Some(path) = self.find_verified_blob(digest)? {
1706 return Ok(path);
1707 }
1708 let _prefetch_permit = match prefetch_limit {
1709 Some(limit) => Some(limit.acquire().await?),
1710 None => None,
1711 };
1712 let _permit = self.remote_transfers.acquire().await?;
1713 let reservation = self.reserve_remote_download(digest.size)?;
1714 self.stats
1715 .remote_blob_requests
1716 .fetch_add(1, Ordering::Relaxed);
1717 let transfer_timer =
1718 AtomicDurationTimer::start(&self.stats.remote_blob_transfer_duration_ns);
1719 let temporary = remote
1720 .get_blob_file(digest, self.remote_staging_dir.as_path())
1721 .await?;
1722 drop(transfer_timer);
1723 let _cas_timer = AtomicDurationTimer::start(&self.stats.local_cas_write_duration_ns);
1724 let path = self.cas.store_verified_file(digest, temporary.path())?;
1725 reservation.commit(digest.size);
1726 self.remember_verified_blob(digest, &path);
1727 self.stats.stores.fetch_add(1, Ordering::Relaxed);
1728 self.stats
1729 .stored_bytes
1730 .fetch_add(digest.size, Ordering::Relaxed);
1731 self.stats
1732 .downloaded_bytes
1733 .fetch_add(digest.size, Ordering::Relaxed);
1734 Ok(path)
1735 }
1736
1737 async fn respond(&self, request: AgentRequest) -> AgentResponse {
1738 let result = match request {
1739 AgentRequest::FindBlob { digest } => self.find_blob(&digest).await,
1740 AgentRequest::FindBlobs { digests } => self.find_blobs(digests).await,
1741 AgentRequest::StoreBlob { digest, source } => self.store_blob(&digest, &source).await,
1742 AgentRequest::FindActionResult { action } => {
1743 self.stats.lookups.fetch_add(1, Ordering::Relaxed);
1744 self.find_action_result(&action).await
1745 }
1746 AgentRequest::RecordActionHit { action, restore } => {
1747 self.record_action_hit(&action, restore)
1748 }
1749 AgentRequest::RecordBypass { kind } => {
1750 *self.stats.bypasses.lock().unwrap().entry(kind).or_insert(0) += 1;
1751 Ok(AgentResponse::BypassRecorded)
1752 }
1753 AgentRequest::RecordUnconsulted => {
1754 self.stats.unconsulted.fetch_add(1, Ordering::Relaxed);
1755 Ok(AgentResponse::UnconsultedRecorded)
1756 }
1757 AgentRequest::RecordCompilerInvocation {
1758 outcome,
1759 crate_name,
1760 duration_ns,
1761 } => self.record_compiler_invocation(&outcome, crate_name.as_deref(), duration_ns),
1762 AgentRequest::RecordActionVerification { matched, restore } => {
1763 self.record_materialization(restore);
1764 self.stats.verifications.fetch_add(1, Ordering::Relaxed);
1765 if !matched {
1766 self.stats.divergences.fetch_add(1, Ordering::Relaxed);
1767 }
1768 Ok(AgentResponse::ActionVerificationRecorded)
1769 }
1770 AgentRequest::StoreActionResult { result } => self.store_action_result(&result).await,
1771 AgentRequest::FindActionPrediction { task, invocation } => {
1772 self.find_action_prediction(&task, &invocation)
1773 }
1774 AgentRequest::RecordActionPrediction { task, prediction } => {
1775 self.record_action_prediction(&task, prediction)
1776 }
1777 AgentRequest::FindExecutableIdentity {
1778 executable,
1779 environment,
1780 } => self.find_executable_identity(executable, environment),
1781 AgentRequest::StoreExecutableIdentity {
1782 executable,
1783 environment,
1784 stdout,
1785 } => self.store_executable_identity(executable, environment, stdout),
1786 AgentRequest::Hello { .. } => {
1787 Err(eyre::eyre!("hello is only valid as the first request"))
1788 }
1789 };
1790 result.unwrap_or_else(|error| AgentResponse::Error {
1791 message: error.to_string(),
1792 })
1793 }
1794
1795 async fn find_blob(&self, digest: &CacheDigest) -> Result<AgentResponse> {
1796 if let Some(path) = self.find_verified_blob(digest)? {
1797 return Ok(AgentResponse::Blob { path: Some(path) });
1798 }
1799 if !self.remote_mode.reads() {
1800 return Ok(AgentResponse::Blob { path: None });
1801 }
1802 let Some(remote) = &self.remote else {
1803 return Ok(AgentResponse::Blob { path: None });
1804 };
1805 match self.fetch_remote_blob(remote, digest).await {
1806 Ok(path) => Ok(AgentResponse::Blob { path: Some(path) }),
1807 Err(error) => {
1808 warn!(
1809 "remote cache blob lookup failed for {}: {error}",
1810 digest.hash
1811 );
1812 Ok(AgentResponse::Blob { path: None })
1813 }
1814 }
1815 }
1816
1817 async fn find_blobs(&self, digests: Vec<CacheDigest>) -> Result<AgentResponse> {
1818 let mut paths = BTreeMap::new();
1819 let mut missing = Vec::new();
1820 for digest in &digests {
1821 match self.find_verified_blob(digest)? {
1822 Some(path) => {
1823 paths.insert(digest.clone(), path);
1824 }
1825 None => {
1826 missing.push(digest.clone());
1827 }
1828 }
1829 }
1830
1831 if !missing.is_empty()
1832 && self.remote_mode.reads()
1833 && let Some(remote) = &self.remote
1834 {
1835 paths.extend(self.fetch_remote_blobs(remote, missing, None).await);
1836 }
1837
1838 Ok(AgentResponse::Blobs {
1839 paths: digests
1840 .into_iter()
1841 .map(|digest| paths.get(&digest).cloned())
1842 .collect(),
1843 })
1844 }
1845
1846 async fn store_blob(&self, digest: &CacheDigest, source: &Path) -> Result<AgentResponse> {
1847 let remote = if self.remote_mode.writes() {
1848 self.remote.as_deref()
1849 } else {
1850 None
1851 };
1852 let path = {
1853 let lock = self.write_lock(digest);
1854 let _guard = lock.lock().await;
1855 if let Some(path) = self.find_verified_blob(digest)? {
1856 path
1857 } else {
1858 let path = self.cas.store_file(digest, source)?;
1859 self.remember_verified_blob(digest, &path);
1860 self.stats.stores.fetch_add(1, Ordering::Relaxed);
1861 self.stats
1862 .stored_bytes
1863 .fetch_add(digest.size, Ordering::Relaxed);
1864 path
1865 }
1866 };
1867 if let Some(remote) = remote {
1868 let _permit = self.remote_transfers.acquire().await?;
1869 if let Err(error) = remote
1870 .put_blob(&BlobUpload {
1871 digest: digest.clone(),
1872 source: BlobSource::Path(path.clone()),
1873 })
1874 .await
1875 {
1876 self.note_remote_failure();
1877 warn!(
1878 "remote cache blob upload failed for {}: {error}",
1879 digest.hash
1880 );
1881 } else {
1882 self.stats
1883 .uploaded_bytes
1884 .fetch_add(digest.size, Ordering::Relaxed);
1885 }
1886 }
1887 Ok(AgentResponse::Stored { path })
1888 }
1889
1890 fn find_verified_blob(&self, digest: &CacheDigest) -> Result<Option<PathBuf>> {
1891 let remembered = self.verified_blobs.lock().unwrap().get(digest).cloned();
1892 if let Some(path) = remembered {
1893 if digest.matches_file(&path).unwrap_or(false) {
1894 return Ok(Some(path));
1895 }
1896 self.verified_blobs.lock().unwrap().remove(digest);
1897 }
1898 let path = self.cas.find(digest)?;
1899 if let Some(path) = &path {
1900 self.remember_verified_blob(digest, path);
1901 }
1902 Ok(path)
1903 }
1904
1905 fn remember_verified_blob(&self, digest: &CacheDigest, path: &Path) {
1906 self.verified_blobs
1907 .lock()
1908 .unwrap()
1909 .insert(digest.clone(), path.to_path_buf());
1910 }
1911
1912 async fn find_action_result(&self, action: &CacheDigest) -> Result<AgentResponse> {
1913 if let Some(result) = self.actions.find(action)? {
1914 return Ok(AgentResponse::ActionResult {
1915 result: Some(result),
1916 });
1917 }
1918 if !self.remote_mode.reads() {
1919 return Ok(AgentResponse::ActionResult { result: None });
1920 }
1921 let Some(remote) = &self.remote else {
1922 return Ok(AgentResponse::ActionResult { result: None });
1923 };
1924 let lock = self.action_lock(action);
1925 let _guard = lock.lock().await;
1926 if let Some(result) = self.actions.find(action)? {
1927 return Ok(AgentResponse::ActionResult {
1928 result: Some(result),
1929 });
1930 }
1931 if let Some(result) = self
1932 .pending_remote_actions
1933 .lock()
1934 .unwrap()
1935 .get(action)
1936 .cloned()
1937 {
1938 return Ok(AgentResponse::ActionResult {
1939 result: Some(result),
1940 });
1941 }
1942 let _permit = self.remote_transfers.acquire().await?;
1943 match self.get_remote_action_result(remote, action).await {
1944 Ok(Some(result)) => {
1945 self.pending_remote_actions
1946 .lock()
1947 .unwrap()
1948 .insert(action.clone(), result.clone());
1949 Ok(AgentResponse::ActionResult {
1950 result: Some(result),
1951 })
1952 }
1953 Ok(None) => Ok(AgentResponse::ActionResult { result: None }),
1954 Err(error) => {
1955 self.note_remote_failure();
1956 warn!(
1957 "remote cache action lookup failed for {}: {error}",
1958 action.hash
1959 );
1960 Ok(AgentResponse::ActionResult { result: None })
1961 }
1962 }
1963 }
1964
1965 async fn store_action_result(&self, result: &RemoteActionResult) -> Result<AgentResponse> {
1966 let path = self.actions.store(result)?;
1967 if self.remote_mode.writes()
1968 && let Some(remote) = &self.remote
1969 {
1970 let _permit = self.remote_transfers.acquire().await?;
1971 if let Err(error) = remote.put_action_result(result).await {
1972 self.note_remote_failure();
1973 warn!(
1974 "remote cache action upload failed for {}: {error}",
1975 result.action.hash
1976 );
1977 }
1978 }
1979 Ok(AgentResponse::ActionStored { path })
1980 }
1981
1982 fn note_remote_failure(&self) {
1984 self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
1985 }
1986
1987 async fn get_remote_action_result(
1988 &self,
1989 remote: &RemoteCacheClient,
1990 action: &CacheDigest,
1991 ) -> Result<Option<RemoteActionResult>> {
1992 self.stats
1993 .remote_action_lookups
1994 .fetch_add(1, Ordering::Relaxed);
1995 let _timer = AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
1996 remote.get_action_result(action).await
1997 }
1998
1999 fn record_action_hit(
2000 &self,
2001 action: &CacheDigest,
2002 restore: RestoreStats,
2003 ) -> Result<AgentResponse> {
2004 if self.actions.find(action)?.is_none() {
2005 let pending = self.pending_remote_actions.lock().unwrap().remove(action);
2006 if let Some(result) = pending {
2007 self.actions.store(&result)?;
2008 } else {
2009 bail!("cannot record a hit for a missing action result");
2010 }
2011 }
2012 self.record_restore(restore);
2013 self.stats.hits.fetch_add(1, Ordering::Relaxed);
2014 Ok(AgentResponse::ActionHitRecorded)
2015 }
2016
2017 fn record_restore(&self, restore: RestoreStats) {
2018 self.record_materialization(restore);
2019 atomic_saturating_add(
2020 &self.stats.avoided_compiler_duration_ns,
2021 restore.avoided_compiler_duration_ns,
2022 );
2023 atomic_saturating_add(&self.stats.restored_output_files, restore.output_files);
2024 atomic_saturating_add(&self.stats.restored_output_bytes, restore.output_bytes);
2025 atomic_saturating_add(
2026 &self.stats.reflinked_output_files,
2027 restore.reflinked_output_files,
2028 );
2029 atomic_saturating_add(
2030 &self.stats.reflinked_output_bytes,
2031 restore.reflinked_output_bytes,
2032 );
2033 atomic_saturating_add(&self.stats.copied_output_files, restore.copied_output_files);
2034 atomic_saturating_add(&self.stats.copied_output_bytes, restore.copied_output_bytes);
2035 }
2036
2037 fn record_compiler_invocation(
2038 &self,
2039 outcome: &str,
2040 crate_name: Option<&str>,
2041 duration_ns: u64,
2042 ) -> Result<AgentResponse> {
2043 if !matches!(outcome, "miss" | "unconsulted" | "bypass" | "verification") {
2044 bail!("invalid compiler invocation outcome");
2045 }
2046 if let Some(crate_name) = crate_name
2047 && (crate_name.len() > 256 || crate_name.contains(['\0', '\n', '\r']))
2048 {
2049 bail!("invalid compiler crate name");
2050 }
2051 let mut compiler = self.stats.compiler.lock().unwrap();
2052 let stats = compiler.entry(outcome.to_string()).or_default();
2053 stats.invocations = stats.invocations.saturating_add(1);
2054 stats.duration_ns = stats.duration_ns.saturating_add(duration_ns);
2055 drop(compiler);
2056 if outcome != "verification"
2057 && let Some(crate_name) = crate_name.filter(|name| !name.is_empty())
2058 {
2059 let mut slow = self.stats.slow_compilations.lock().unwrap();
2060 let duration = slow.entry(crate_name.to_string()).or_default();
2061 *duration = duration.saturating_add(duration_ns);
2062 }
2063 Ok(AgentResponse::CompilerInvocationRecorded)
2064 }
2065
2066 fn record_materialization(&self, restore: RestoreStats) {
2067 atomic_saturating_add(&self.stats.materialization_duration_ns, restore.duration_ns);
2068 }
2069
2070 fn find_action_prediction(
2071 &self,
2072 task: &str,
2073 invocation: &CacheDigest,
2074 ) -> Result<AgentResponse> {
2075 validate_task_identity(task)?;
2076 invocation.validate()?;
2077 let prediction = self
2078 .task_actions
2079 .lock()
2080 .unwrap()
2081 .get(task)
2082 .and_then(|state| state.predictions.get(invocation))
2083 .cloned();
2084 Ok(AgentResponse::ActionPrediction { prediction })
2085 }
2086
2087 fn record_action_prediction(
2088 &self,
2089 task: &str,
2090 prediction: ActionPrediction,
2091 ) -> Result<AgentResponse> {
2092 validate_task_identity(task)?;
2093 validate_action_prediction(&prediction)?;
2094 let mut tasks = self.task_actions.lock().unwrap();
2095 let state = tasks.entry(task.to_string()).or_default();
2096 if !state.predictions.contains_key(&prediction.invocation)
2097 && state.predictions.len() >= MAX_TASK_ACTION_PREDICTIONS
2098 {
2099 bail!("task action manifest contains too many predictions");
2100 }
2101 state
2102 .predictions
2103 .insert(prediction.invocation.clone(), prediction.clone());
2104 state
2105 .pending_predictions
2106 .insert(prediction.invocation.clone(), prediction);
2107 Ok(AgentResponse::ActionPredictionRecorded)
2108 }
2109
2110 fn executable_identity_key(
2111 &self,
2112 executable: PathBuf,
2113 environment: BTreeMap<String, Option<String>>,
2114 ) -> Result<ExecutableIdentityKey> {
2115 if !environment
2116 .keys()
2117 .all(|name| matches!(name.as_str(), "RUSTUP_HOME" | "RUSTUP_TOOLCHAIN"))
2118 {
2119 bail!("executable identity contains an unsupported environment variable");
2120 }
2121 Ok(ExecutableIdentityKey {
2122 executable,
2123 environment,
2124 })
2125 }
2126
2127 fn find_executable_identity(
2128 &self,
2129 executable: PathBuf,
2130 environment: BTreeMap<String, Option<String>>,
2131 ) -> Result<AgentResponse> {
2132 let key = self.executable_identity_key(executable, environment)?;
2133 let stdout = self
2134 .executable_identities
2135 .lock()
2136 .unwrap()
2137 .get(&key)
2138 .cloned();
2139 Ok(AgentResponse::ExecutableIdentity { stdout })
2140 }
2141
2142 fn store_executable_identity(
2143 &self,
2144 executable: PathBuf,
2145 environment: BTreeMap<String, Option<String>>,
2146 stdout: Vec<u8>,
2147 ) -> Result<AgentResponse> {
2148 if stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE {
2149 bail!("executable identity exceeds {MAX_EXECUTABLE_IDENTITY_SIZE} bytes");
2150 }
2151 let key = self.executable_identity_key(executable, environment)?;
2152 let mut identities = self.executable_identities.lock().unwrap();
2153 let is_new = !identities.contains_key(&key);
2154 let previous_size = identities.get(&key).map_or(0, Vec::len);
2155 if is_new && identities.len() >= MAX_EXECUTABLE_IDENTITIES {
2156 bail!("executable identity cache contains too many entries");
2157 }
2158 let retained_bytes = identities.values().map(Vec::len).sum::<usize>();
2159 if retained_bytes - previous_size + stdout.len() > MAX_EXECUTABLE_IDENTITY_BYTES {
2160 bail!("executable identity cache contains too many bytes");
2161 }
2162 identities.insert(key, stdout.clone());
2163 Ok(AgentResponse::ExecutableIdentity {
2164 stdout: Some(stdout),
2165 })
2166 }
2167
2168 pub async fn handle_connection<S>(&self, stream: S) -> Result<()>
2170 where
2171 S: AsyncRead + AsyncWrite + Unpin,
2172 {
2173 let (reader, mut writer) = tokio::io::split(stream);
2174 let mut reader = BufReader::new(reader);
2175 let hello = read_request(&mut reader)
2176 .await?
2177 .ok_or_else(|| eyre::eyre!("connection closed before the agent handshake"))?;
2178 let request: AgentRequest = serde_json::from_str(&hello)?;
2179 match request {
2180 AgentRequest::Hello {
2181 protocol,
2182 client_version,
2183 } if protocol == AGENT_PROTOCOL_VERSION && client_version == self.version.as_ref() => {}
2184 AgentRequest::Hello { protocol, .. } if protocol != AGENT_PROTOCOL_VERSION => {
2185 send_response(
2186 &mut writer,
2187 &AgentResponse::Error {
2188 message: format!(
2189 "unsupported agent protocol {protocol}; expected {AGENT_PROTOCOL_VERSION}"
2190 ),
2191 },
2192 )
2193 .await?;
2194 return Ok(());
2195 }
2196 AgentRequest::Hello { client_version, .. } => {
2197 send_response(
2198 &mut writer,
2199 &AgentResponse::Error {
2200 message: format!(
2201 "cache client {client_version} does not match agent {}",
2202 self.version
2203 ),
2204 },
2205 )
2206 .await?;
2207 return Ok(());
2208 }
2209 _ => bail!("the first agent request must be hello"),
2210 }
2211 send_response(
2212 &mut writer,
2213 &AgentResponse::Hello {
2214 protocol: AGENT_PROTOCOL_VERSION,
2215 agent_version: self.version.to_string(),
2216 },
2217 )
2218 .await?;
2219
2220 while let Some(line) = read_request(&mut reader).await? {
2221 let response = match serde_json::from_str(&line) {
2222 Ok(request) => self.respond(request).await,
2223 Err(error) => AgentResponse::Error {
2224 message: format!("invalid agent request: {error}"),
2225 },
2226 };
2227 send_response(&mut writer, &response).await?;
2228 }
2229 Ok(())
2230 }
2231}
2232
2233async fn read_request<R>(reader: &mut R) -> Result<Option<String>>
2239where
2240 R: AsyncBufRead + Unpin,
2241{
2242 let mut line = Vec::new();
2243 loop {
2244 let available = reader.fill_buf().await?;
2245 if available.is_empty() {
2246 break;
2247 }
2248 let (consumed, complete) = match available.iter().position(|byte| *byte == b'\n') {
2249 Some(index) => (index, true),
2250 None => (available.len(), false),
2251 };
2252 if line.len() + consumed > MAX_REQUEST_BYTES {
2253 bail!("agent request exceeded {MAX_REQUEST_BYTES} bytes");
2254 }
2255 line.extend_from_slice(&available[..consumed]);
2256 reader.consume(consumed + usize::from(complete));
2258 if complete {
2259 return Ok(Some(String::from_utf8(line)?));
2260 }
2261 }
2262 if line.is_empty() {
2263 Ok(None)
2264 } else {
2265 Ok(Some(String::from_utf8(line)?))
2266 }
2267}
2268
2269pub fn is_task_identity(task: &str) -> bool {
2275 task.len() == 64
2276 && task
2277 .bytes()
2278 .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
2279}
2280
2281fn validate_task_identity(task: &str) -> Result<()> {
2282 if !is_task_identity(task) {
2283 bail!("invalid task action identity");
2284 }
2285 Ok(())
2286}
2287
2288fn task_manifest_dir(store: &Path) -> PathBuf {
2290 store.join("task-manifests").join("v1")
2291}
2292
2293pub fn task_manifest_actions(store: &Path, task: &str) -> Result<Vec<CacheDigest>> {
2301 validate_task_identity(task)?;
2302 let path = task_manifest_dir(store).join(format!("{task}.json"));
2303 let bytes = match fs::read(&path) {
2304 Ok(bytes) => bytes,
2305 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
2306 Err(error) => {
2307 return Err(error).wrap_err_with(|| format!("failed to read {}", path.display()));
2308 }
2309 };
2310 let Ok(manifest) = serde_json::from_slice::<TaskActionManifest>(&bytes) else {
2311 return Ok(Vec::new());
2312 };
2313 if validate_task_manifest(&manifest, task).is_err() {
2314 return Ok(Vec::new());
2315 }
2316 Ok(manifest
2317 .predictions
2318 .into_iter()
2319 .map(|prediction| prediction.action)
2320 .collect())
2321}
2322
2323fn validate_action_prediction(prediction: &ActionPrediction) -> Result<()> {
2324 if prediction.validate() {
2325 Ok(())
2326 } else {
2327 bail!("invalid action prediction")
2328 }
2329}
2330
2331fn validate_task_manifest(manifest: &TaskActionManifest, task: &str) -> Result<()> {
2332 if manifest.task == task && manifest.validate() {
2333 Ok(())
2334 } else {
2335 bail!("invalid task action manifest")
2336 }
2337}
2338
2339fn merge_task_manifests(
2340 task: &str,
2341 base: Option<TaskActionManifest>,
2342 update: TaskActionManifest,
2343) -> Result<TaskActionManifest> {
2344 validate_task_manifest(&update, task)?;
2345 let mut predictions = BTreeMap::new();
2346 if let Some(base) = base {
2347 validate_task_manifest(&base, task)?;
2348 predictions.extend(
2349 base.predictions
2350 .into_iter()
2351 .map(|prediction| (prediction.invocation.clone(), prediction)),
2352 );
2353 }
2354 predictions.extend(
2355 update
2356 .predictions
2357 .into_iter()
2358 .map(|prediction| (prediction.invocation.clone(), prediction)),
2359 );
2360 let manifest = TaskActionManifest {
2361 version: TASK_ACTION_MANIFEST_VERSION,
2362 task: task.to_owned(),
2363 predictions: predictions.into_values().collect(),
2364 };
2365 validate_task_manifest(&manifest, task)?;
2366 Ok(manifest)
2367}
2368
2369fn merge_remote_task_manifest(
2370 task: &str,
2371 remote: TaskActionManifest,
2372 local: TaskActionManifest,
2373) -> (TaskActionManifest, bool) {
2374 match merge_task_manifests(task, Some(remote), local.clone()) {
2375 Ok(manifest) => (manifest, true),
2376 Err(error) => {
2377 warn!("remote task action manifest merge failed for {task}: {error}");
2378 (local, false)
2379 }
2380 }
2381}
2382
2383async fn send_response(
2384 writer: &mut (impl AsyncWrite + Unpin),
2385 response: &AgentResponse,
2386) -> Result<()> {
2387 let mut encoded = serde_json::to_vec(response)?;
2388 encoded.push(b'\n');
2389 writer.write_all(&encoded).await?;
2390 writer.flush().await?;
2391 Ok(())
2392}
2393
2394#[cfg(test)]
2395#[path = "agent_tests.rs"]
2396mod tests;