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