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