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