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 self.release_bypass(&mut connection);
1883 responses
1884 }
1885
1886 fn release_bypass(&self, connection: &mut ConnectionUploads) {
1888 if let Some(kind) = connection.take_bypass() {
1889 self.emit(|| AgentEvent::Bypass { kind });
1890 }
1891 }
1892
1893 fn write_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1894 Self::digest_lock(&self.write_locks, digest)
1895 }
1896
1897 fn action_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1898 Self::digest_lock(&self.action_locks, digest)
1899 }
1900
1901 fn digest_lock(
1902 locks: &Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>,
1903 digest: &CacheDigest,
1904 ) -> Arc<tokio::sync::Mutex<()>> {
1905 let mut locks = locks.lock().unwrap();
1906 locks.retain(|_, lock| lock.strong_count() > 0);
1907 if let Some(lock) = locks.get(digest).and_then(Weak::upgrade) {
1908 return lock;
1909 }
1910 let lock = Arc::new(tokio::sync::Mutex::new(()));
1911 locks.insert(digest.clone(), Arc::downgrade(&lock));
1912 lock
1913 }
1914
1915 #[cfg(test)]
1920 async fn respond(&self, request: AgentRequest) -> AgentResponse {
1921 let mut connection = ConnectionUploads::default();
1922 let response = self.respond_on(request, &mut connection).await;
1923 self.release_bypass(&mut connection);
1924 response
1925 }
1926
1927 async fn respond_on(
1928 &self,
1929 request: AgentRequest,
1930 connection: &mut ConnectionUploads,
1931 ) -> AgentResponse {
1932 match &request {
1933 AgentRequest::FindActionPrediction { task, .. }
1934 | AgentRequest::RecordActionPrediction { task, .. } => {
1935 self.ensure_task_loaded(task).await;
1936 }
1937 _ => {}
1938 }
1939 let result = match request {
1940 AgentRequest::BeginTask { task } => self
1941 .begin_task(&task)
1942 .await
1943 .map(|run| AgentResponse::TaskBegun { run }),
1944 AgentRequest::CommitTask { run } => self
1945 .commit_task(&run)
1946 .await
1947 .map(|()| AgentResponse::TaskCommitted),
1948 AgentRequest::FindBlob { digest } => self.find_blob(&digest).await,
1949 AgentRequest::FindBlobs { digests } => self.find_blobs(digests).await,
1950 AgentRequest::StoreBlob { digest, source } => {
1951 self.store_blob(&digest, &source, connection).await
1952 }
1953 AgentRequest::FindActionResult { action } => self.find_action_result(&action).await,
1960 AgentRequest::RecordActionHit {
1961 action,
1962 restore,
1963 crate_name,
1964 } => {
1965 let diagnostic = connection
1966 .take_action_diagnostic("hit", crate_name.as_deref())
1967 .map(|diagnostic| AgentEvent::ActionDiagnostic {
1968 outcome: "hit".into(),
1969 crate_name: crate_name.clone(),
1970 diagnostic,
1971 });
1972 self.record_action_hit(&action, restore, crate_name, diagnostic)
1973 }
1974 AgentRequest::RecordBypass { kind } => {
1975 *self
1976 .stats
1977 .bypasses
1978 .lock()
1979 .unwrap()
1980 .entry(kind.clone())
1981 .or_insert(0) += 1;
1982 if let Some(unclaimed) = connection.hold_bypass(kind) {
1986 self.emit(|| AgentEvent::Bypass { kind: unclaimed });
1987 }
1988 Ok(AgentResponse::BypassRecorded)
1989 }
1990 AgentRequest::RecordUnconsulted => {
1991 self.stats.unconsulted.fetch_add(1, Ordering::Relaxed);
1992 self.emit(|| AgentEvent::Unconsulted);
1993 Ok(AgentResponse::UnconsultedRecorded)
1994 }
1995 AgentRequest::RecordWarning { message } => match parse_action_diagnostic(&message) {
1996 Some(Ok((outcome, crate_name, diagnostic))) => {
1997 connection.record_action_diagnostic(outcome, crate_name, diagnostic);
1998 Ok(AgentResponse::WarningRecorded)
1999 }
2000 Some(Err(error)) => Err(error),
2001 None => self.record_diagnostic(DiagnosticLevel::Warning, message),
2002 },
2003 AgentRequest::RecordDebug { target, message } => (|| {
2004 if target.is_empty()
2005 || target.len() > 256
2006 || target.contains(['\n', '\r', '\0'])
2007 || message.is_empty()
2008 || message.len() > MAX_DIAGNOSTIC_BYTES
2009 || message.contains(['\n', '\r', '\0'])
2010 {
2011 bail!("invalid shim debug record");
2012 }
2013 if target == "mbx::unit-outcome" {
2014 let (unit, outcome): (String, String) = serde_json::from_str(&message)?;
2015 let (outcome, _) = lookup_outcome(&outcome);
2019 if unit.len() > 256
2020 || !unit.contains(':')
2021 || !matches!(
2022 outcome,
2023 "hit" | "miss" | "bypass" | "unconsulted" | "verification"
2024 )
2025 {
2026 bail!("invalid unit outcome");
2027 }
2028 self.stats
2029 .unit_outcomes
2030 .lock()
2031 .unwrap()
2032 .entry(unit)
2033 .or_default()
2034 .insert(outcome.to_string());
2035 } else {
2036 debug!(target: &target, "{message}");
2037 }
2038 Ok(AgentResponse::DebugRecorded)
2039 })(),
2040 AgentRequest::RecordError { message } => {
2041 self.record_diagnostic(DiagnosticLevel::Error, message)
2042 }
2043 AgentRequest::FindFileDigests { scope, files } => self.find_file_digests(scope, files),
2044 AgentRequest::JoinActionPromise {
2045 adapter,
2046 invocation,
2047 } => self.join_action_promise(&adapter, &invocation).await,
2048 AgentRequest::CompleteActionPromise { claim, prediction } => {
2049 self.complete_action_promise(&claim, &prediction).await
2050 }
2051 AgentRequest::ResolveFileDigests { scope, files } => {
2052 self.resolve_file_digests(scope, files).await
2053 }
2054 AgentRequest::RecordFileDigests { scope, entries } => {
2055 self.record_file_digests(scope, entries)
2056 }
2057 AgentRequest::RecordWrapperTiming { timing } => (|| {
2058 if timing.spans.len() > 512
2059 || timing.phases_ns.len() > 32
2060 || timing.adapter.len() > 64
2061 || timing.unit.as_ref().is_some_and(|s| s.len() > 1024)
2062 || timing.unit_id.as_ref().is_some_and(|s| s.len() > 256)
2063 || timing.dependencies.len() > 4096
2064 || timing.dependencies.iter().any(|s| s.len() > 256)
2065 || timing.phases_ns.keys().any(|s| s.len() > 64)
2066 || timing.spans.iter().any(|s| s.name.len() > 64)
2067 {
2068 bail!("wrapper timing exceeds telemetry limits");
2069 }
2070 {
2071 let mut phases = self.stats.wrapper_phases_ns.lock().unwrap();
2072 for (phase, duration) in &timing.phases_ns {
2073 let total = phases.entry(phase.clone()).or_default();
2074 *total = total.saturating_add(*duration);
2075 }
2076 }
2077 self.emit(|| AgentEvent::WrapperTiming { timing });
2078 Ok(AgentResponse::WrapperTimingRecorded)
2079 })(),
2080 AgentRequest::RecordCompilerInvocation {
2081 outcome,
2082 crate_name,
2083 duration_ns,
2084 } => {
2085 let (outcome, incremental) = lookup_outcome(&outcome);
2089 let outcome = outcome.to_string();
2090 let bypass = if outcome == "bypass" {
2091 connection.take_bypass()
2092 } else {
2093 self.release_bypass(connection);
2094 None
2095 };
2096 let diagnostic = connection
2097 .take_action_diagnostic(&outcome, crate_name.as_deref())
2098 .map(|diagnostic| AgentEvent::ActionDiagnostic {
2099 outcome: outcome.clone(),
2100 crate_name: crate_name.clone(),
2101 diagnostic,
2102 });
2103 self.record_compiler_invocation(
2104 &outcome,
2105 incremental,
2106 crate_name.as_deref(),
2107 duration_ns,
2108 diagnostic,
2109 bypass,
2110 )
2111 }
2112 AgentRequest::RecordActionVerification { matched, restore } => {
2113 self.record_materialization(restore);
2114 self.stats.verifications.fetch_add(1, Ordering::Relaxed);
2115 if !matched {
2116 self.stats.divergences.fetch_add(1, Ordering::Relaxed);
2117 }
2118 self.emit(|| AgentEvent::Verification { matched, restore });
2119 Ok(AgentResponse::ActionVerificationRecorded)
2120 }
2121 AgentRequest::StoreActionResult { result } => {
2122 self.store_action_result(&result, connection).await
2123 }
2124 AgentRequest::FindActionPrediction { task, invocation } => {
2125 self.find_action_prediction(&task, &invocation)
2126 }
2127 AgentRequest::RecordActionPrediction { task, prediction } => {
2128 self.record_action_prediction(&task, prediction)
2129 }
2130 AgentRequest::FindExecutableIdentity {
2131 executable,
2132 environment,
2133 } => self.find_executable_identity(executable, environment),
2134 AgentRequest::StoreExecutableIdentity {
2135 executable,
2136 environment,
2137 stdout,
2138 pins,
2139 } => self.store_executable_identity(executable, environment, stdout, pins),
2140 AgentRequest::Hello { .. } => {
2141 Err(eyre::eyre!("hello is only valid as the first request"))
2142 }
2143 };
2144 result.unwrap_or_else(|error| AgentResponse::Error {
2145 message: error.to_string(),
2146 })
2147 }
2148
2149 async fn find_blob(&self, digest: &CacheDigest) -> Result<AgentResponse> {
2150 if let Some(path) = self.find_verified_blob(digest)? {
2151 return Ok(AgentResponse::Blob { path: Some(path) });
2152 }
2153 if !self.remote_mode.reads() {
2154 return Ok(AgentResponse::Blob { path: None });
2155 }
2156 let Some(remote) = &self.remote else {
2157 return Ok(AgentResponse::Blob { path: None });
2158 };
2159 match self.fetch_remote_blob(remote, digest).await {
2160 Ok(path) => Ok(AgentResponse::Blob { path: Some(path) }),
2161 Err(error) => {
2162 warn!(
2163 "remote cache blob lookup failed for {}: {error}",
2164 digest.hash
2165 );
2166 Ok(AgentResponse::Blob { path: None })
2167 }
2168 }
2169 }
2170
2171 async fn find_blobs(&self, digests: Vec<CacheDigest>) -> Result<AgentResponse> {
2172 let mut paths = BTreeMap::new();
2173 let mut missing = Vec::new();
2174 for digest in &digests {
2175 match self.find_verified_blob(digest)? {
2176 Some(path) => {
2177 paths.insert(digest.clone(), path);
2178 }
2179 None => {
2180 missing.push(digest.clone());
2181 }
2182 }
2183 }
2184
2185 if !missing.is_empty()
2186 && self.remote_mode.reads()
2187 && let Some(remote) = &self.remote
2188 {
2189 paths.extend(self.fetch_remote_blobs(remote, missing, None).await);
2190 }
2191
2192 Ok(AgentResponse::Blobs {
2193 paths: digests
2194 .into_iter()
2195 .map(|digest| paths.get(&digest).cloned())
2196 .collect(),
2197 })
2198 }
2199
2200 async fn store_blob(
2201 &self,
2202 digest: &CacheDigest,
2203 source: &Path,
2204 connection: &mut ConnectionUploads,
2205 ) -> Result<AgentResponse> {
2206 let path = {
2207 let lock = self.write_lock(digest);
2208 let _guard = lock.lock().await;
2209 if let Some(path) = self.find_verified_blob(digest)? {
2210 path
2211 } else {
2212 let timer = AtomicDurationTimer::start(&self.stats.local_cas_write_duration_ns);
2213 let path = self.cas.store_file(digest, source)?;
2214 drop(timer);
2215 self.remember_verified_blob(digest, &path);
2216 self.stats.stores.fetch_add(1, Ordering::Relaxed);
2217 self.stats
2218 .stored_bytes
2219 .fetch_add(digest.size, Ordering::Relaxed);
2220 path
2221 }
2222 };
2223 if let Some(uploads) = &self.uploads {
2226 uploads.queue_blob(digest, path.clone(), connection);
2227 }
2228 Ok(AgentResponse::Stored { path })
2229 }
2230
2231 fn find_verified_blob(&self, digest: &CacheDigest) -> Result<Option<PathBuf>> {
2232 let remembered = self.verified_blobs.lock().unwrap().get(digest).cloned();
2233 if let Some(remembered) = remembered {
2234 if remembered.is_unchanged() {
2245 return Ok(Some(remembered.path));
2246 }
2247 self.verified_blobs.lock().unwrap().remove(digest);
2248 }
2249 let path = match self.cas.find(digest) {
2256 Ok(path) => path,
2257 Err(error) => {
2258 warn!(
2259 "local CAS blob {} is not restorable and will be republished: {error}",
2260 digest.hash
2261 );
2262 return Ok(None);
2263 }
2264 };
2265 if let Some(path) = &path {
2266 self.remember_verified_blob(digest, path);
2267 }
2268 Ok(path)
2269 }
2270
2271 fn remember_verified_blob(&self, digest: &CacheDigest, path: &Path) {
2272 let Some(verified) = VerifiedBlob::describe(path) else {
2273 return;
2277 };
2278 self.verified_blobs
2279 .lock()
2280 .unwrap()
2281 .insert(digest.clone(), verified);
2282 }
2283
2284 async fn find_action_result(&self, action: &CacheDigest) -> Result<AgentResponse> {
2285 if let Some(result) = self.actions.find(action)? {
2286 return Ok(AgentResponse::ActionResult {
2287 result: Some(result),
2288 });
2289 }
2290 if !self.remote_mode.reads() {
2291 return Ok(AgentResponse::ActionResult { result: None });
2292 }
2293 let Some(remote) = &self.remote else {
2294 return Ok(AgentResponse::ActionResult { result: None });
2295 };
2296 let lock = self.action_lock(action);
2297 let _guard = lock.lock().await;
2298 if let Some(result) = self.actions.find(action)? {
2299 return Ok(AgentResponse::ActionResult {
2300 result: Some(result),
2301 });
2302 }
2303 if let Some(result) = self
2304 .pending_remote_actions
2305 .lock()
2306 .unwrap()
2307 .get(action)
2308 .cloned()
2309 {
2310 return Ok(AgentResponse::ActionResult {
2311 result: Some(result),
2312 });
2313 }
2314 let _permit = self.remote_transfers.acquire().await?;
2315 match self.get_remote_action_result(remote, action).await {
2316 Ok(Some(result)) => {
2317 self.pending_remote_actions
2318 .lock()
2319 .unwrap()
2320 .insert(action.clone(), result.clone());
2321 Ok(AgentResponse::ActionResult {
2322 result: Some(result),
2323 })
2324 }
2325 Ok(None) => Ok(AgentResponse::ActionResult { result: None }),
2326 Err(error) => {
2327 self.note_remote_failure();
2328 warn!(
2329 "remote cache action lookup failed for {}: {error}",
2330 action.hash
2331 );
2332 Ok(AgentResponse::ActionResult { result: None })
2333 }
2334 }
2335 }
2336
2337 async fn store_action_result(
2338 &self,
2339 result: &RemoteActionResult,
2340 connection: &ConnectionUploads,
2341 ) -> Result<AgentResponse> {
2342 let path = self.actions.store(result)?;
2343 if let Some(uploads) = &self.uploads {
2344 uploads.queue_action_result(result, connection);
2345 }
2346 Ok(AgentResponse::ActionStored { path })
2347 }
2348
2349 async fn join_action_promise(
2350 &self,
2351 adapter: &str,
2352 invocation: &CacheDigest,
2353 ) -> Result<AgentResponse> {
2354 if !self.remote_mode.reads() || !self.remote_mode.writes() {
2355 return Ok(AgentResponse::ActionPromise {
2356 claim: None,
2357 prediction: None,
2358 });
2359 }
2360 let Some(remote) = &self.remote else {
2361 return Ok(AgentResponse::ActionPromise {
2362 claim: None,
2363 prediction: None,
2364 });
2365 };
2366 let deadline = Instant::now() + ACTION_PROMISE_WAIT;
2367 loop {
2368 let _permit = self.remote_transfers.acquire().await?;
2369 let state = remote.join_action_promise(invocation, adapter).await;
2370 drop(_permit);
2371 match state {
2372 Ok(Some(ActionPromiseState::Claimed { claim })) => {
2373 return Ok(AgentResponse::ActionPromise {
2374 claim: Some(claim),
2375 prediction: None,
2376 });
2377 }
2378 Ok(Some(ActionPromiseState::Complete { prediction })) => {
2379 return Ok(AgentResponse::ActionPromise {
2380 claim: None,
2381 prediction: Some(prediction),
2382 });
2383 }
2384 Ok(Some(ActionPromiseState::Pending { retry_after_ms }))
2385 if Instant::now() < deadline =>
2386 {
2387 tokio::time::sleep(Duration::from_millis(retry_after_ms.clamp(10, 5_000)))
2388 .await;
2389 }
2390 Ok(Some(ActionPromiseState::Pending { .. }) | None) => {
2391 return Ok(AgentResponse::ActionPromise {
2392 claim: None,
2393 prediction: None,
2394 });
2395 }
2396 Ok(Some(_)) => {
2397 return Ok(AgentResponse::ActionPromise {
2398 claim: None,
2399 prediction: None,
2400 });
2401 }
2402 Err(error) => {
2403 self.note_remote_failure();
2404 warn!(
2405 "remote cache action promise failed for {}: {error}",
2406 invocation.hash
2407 );
2408 return Ok(AgentResponse::ActionPromise {
2409 claim: None,
2410 prediction: None,
2411 });
2412 }
2413 }
2414 }
2415 }
2416
2417 async fn complete_action_promise(
2418 &self,
2419 claim: &str,
2420 prediction: &ActionPrediction,
2421 ) -> Result<AgentResponse> {
2422 prediction.validate()?;
2423 let Some(remote) = &self.remote else {
2424 return Ok(AgentResponse::ActionPromiseCompleted);
2425 };
2426 let Some(uploads) = &self.uploads else {
2427 return Ok(AgentResponse::ActionPromiseCompleted);
2428 };
2429 if uploads
2430 .wait_for_actions(std::slice::from_ref(&prediction.action))
2431 .await
2432 .contains(&prediction.action)
2433 {
2434 return Ok(AgentResponse::ActionPromiseCompleted);
2437 }
2438 let completion = ActionPromiseCompletion {
2439 claim: claim.to_string(),
2440 prediction: prediction.clone(),
2441 };
2442 let _permit = self.remote_transfers.acquire().await?;
2443 if let Err(error) = remote
2444 .complete_action_promise(&prediction.invocation, &completion)
2445 .await
2446 {
2447 self.note_remote_failure();
2448 warn!(
2449 "remote cache action promise completion failed for {}: {error}",
2450 prediction.invocation.hash
2451 );
2452 }
2453 Ok(AgentResponse::ActionPromiseCompleted)
2454 }
2455
2456 fn note_remote_failure(&self) {
2458 self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
2459 }
2460
2461 async fn get_remote_action_result(
2462 &self,
2463 remote: &RemoteCacheClient,
2464 action: &CacheDigest,
2465 ) -> Result<Option<RemoteActionResult>> {
2466 self.stats
2467 .remote_action_lookups
2468 .fetch_add(1, Ordering::Relaxed);
2469 let _timer = AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
2470 remote.get_action_result(action).await
2471 }
2472
2473 fn record_action_hit(
2474 &self,
2475 action: &CacheDigest,
2476 restore: RestoreStats,
2477 crate_name: Option<String>,
2478 diagnostic: Option<AgentEvent>,
2479 ) -> Result<AgentResponse> {
2480 validate_crate_name(crate_name.as_deref())?;
2481 if self.actions.find(action)?.is_none() {
2482 let pending = self.pending_remote_actions.lock().unwrap().remove(action);
2483 if let Some(result) = pending {
2484 self.actions.store(&result)?;
2485 } else {
2486 bail!("cannot record a hit for a missing action result");
2487 }
2488 }
2489 self.record_restore(restore);
2490 self.stats.lookups.fetch_add(1, Ordering::Relaxed);
2491 self.stats.hits.fetch_add(1, Ordering::Relaxed);
2492 self.emit_action(
2493 diagnostic,
2494 AgentEvent::ActionHit {
2495 crate_name,
2496 restore,
2497 },
2498 );
2499 Ok(AgentResponse::ActionHitRecorded)
2500 }
2501
2502 fn record_restore(&self, restore: RestoreStats) {
2503 self.record_materialization(restore);
2504 atomic_saturating_add(
2505 &self.stats.avoided_compiler_duration_ns,
2506 restore.avoided_compiler_duration_ns,
2507 );
2508 atomic_saturating_add(&self.stats.restored_output_files, restore.output_files);
2509 atomic_saturating_add(&self.stats.restored_output_bytes, restore.output_bytes);
2510 atomic_saturating_add(
2511 &self.stats.reflinked_output_files,
2512 restore.reflinked_output_files,
2513 );
2514 atomic_saturating_add(
2515 &self.stats.reflinked_output_bytes,
2516 restore.reflinked_output_bytes,
2517 );
2518 atomic_saturating_add(
2519 &self.stats.hardlinked_output_files,
2520 restore.hardlinked_output_files,
2521 );
2522 atomic_saturating_add(
2523 &self.stats.hardlinked_output_bytes,
2524 restore.hardlinked_output_bytes,
2525 );
2526 atomic_saturating_add(&self.stats.copied_output_files, restore.copied_output_files);
2527 atomic_saturating_add(&self.stats.copied_output_bytes, restore.copied_output_bytes);
2528 atomic_saturating_add(&self.stats.reused_output_files, restore.reused_output_files);
2529 atomic_saturating_add(&self.stats.reused_output_bytes, restore.reused_output_bytes);
2530 }
2531
2532 fn record_compiler_invocation(
2533 &self,
2534 outcome: &str,
2535 incremental: bool,
2536 crate_name: Option<&str>,
2537 duration_ns: u64,
2538 diagnostic: Option<AgentEvent>,
2539 bypass: Option<String>,
2540 ) -> Result<AgentResponse> {
2541 if !matches!(outcome, "miss" | "unconsulted" | "bypass" | "verification") {
2542 bail!("invalid compiler invocation outcome");
2543 }
2544 validate_crate_name(crate_name)?;
2545 if incremental {
2546 self.stats
2547 .incremental_compilations
2548 .fetch_add(1, Ordering::Relaxed);
2549 }
2550 if matches!(outcome, "miss" | "verification") {
2552 self.stats.lookups.fetch_add(1, Ordering::Relaxed);
2553 }
2554 let mut compiler = self.stats.compiler.lock().unwrap();
2555 let stats = compiler.entry(outcome.to_string()).or_default();
2556 stats.invocations = stats.invocations.saturating_add(1);
2557 stats.duration_ns = stats.duration_ns.saturating_add(duration_ns);
2558 drop(compiler);
2559 if outcome != "verification"
2560 && let Some(crate_name) = crate_name.filter(|name| !name.is_empty())
2561 {
2562 let mut slow = self.stats.slow_compilations.lock().unwrap();
2563 let duration = slow.entry(crate_name.to_string()).or_default();
2564 *duration = duration.saturating_add(duration_ns);
2565 }
2566 let action = match bypass {
2567 Some(kind) => AgentEvent::BypassedCompilation {
2568 kind,
2569 crate_name: crate_name.map(str::to_string),
2570 duration_ns,
2571 },
2572 None => AgentEvent::CompilerInvocation {
2573 outcome: outcome.to_string(),
2574 crate_name: crate_name.map(str::to_string),
2575 duration_ns,
2576 },
2577 };
2578 self.emit_action(diagnostic, action);
2579 Ok(AgentResponse::CompilerInvocationRecorded)
2580 }
2581
2582 fn record_materialization(&self, restore: RestoreStats) {
2583 atomic_saturating_add(&self.stats.materialization_duration_ns, restore.duration_ns);
2584 }
2585
2586 fn record_diagnostic(&self, level: DiagnosticLevel, message: String) -> Result<AgentResponse> {
2597 if message.is_empty()
2598 || message.len() > MAX_DIAGNOSTIC_BYTES
2599 || message.contains(['\n', '\r', '\0'])
2600 {
2601 bail!("invalid shim {}", level.label());
2602 }
2603 let mut diagnostics = self.diagnostics.lock().unwrap();
2604 let key = (level, message.clone());
2605 if !diagnostics.contains(&key) {
2606 if diagnostics.len() >= MAX_DIAGNOSTICS {
2610 bail!("shim diagnostic limit reached");
2611 }
2612 eprintln!("mbx[{}]: {message}", level.label());
2613 self.emit(|| match level {
2614 DiagnosticLevel::Warning => AgentEvent::Warning {
2615 message: message.clone(),
2616 },
2617 DiagnosticLevel::Error => AgentEvent::Error {
2618 message: message.clone(),
2619 },
2620 });
2621 diagnostics.insert(key);
2622 }
2623 Ok(match level {
2624 DiagnosticLevel::Warning => AgentResponse::WarningRecorded,
2625 DiagnosticLevel::Error => AgentResponse::ErrorRecorded,
2626 })
2627 }
2628
2629 fn find_file_digests(
2636 &self,
2637 scope: FileDigestScope,
2638 files: Vec<FileIdentity>,
2639 ) -> Result<AgentResponse> {
2640 if files.len() > MAX_FILE_DIGEST_BATCH {
2641 bail!("too many file-digest lookups in one request");
2642 }
2643 let ledger = self.file_digests.lock().unwrap();
2644 let digests = files
2645 .into_iter()
2646 .map(|file| {
2647 let recorded = ledger.get(&(scope, file.path.clone()))?;
2648 (recorded.file == file).then(|| recorded.digest.clone())
2649 })
2650 .collect();
2651 Ok(AgentResponse::FileDigests { digests })
2652 }
2653
2654 async fn resolve_file_digests(
2657 &self,
2658 scope: FileDigestScope,
2659 files: Vec<FileIdentity>,
2660 ) -> Result<AgentResponse> {
2661 if files.len() > MAX_FILE_DIGEST_BATCH {
2662 bail!("too many file-digest resolutions in one request");
2663 }
2664 for file in &files {
2665 if !file.path.is_absolute() {
2666 bail!("file-digest resolutions need absolute paths");
2667 }
2668 }
2669 let resolutions = stream::iter(
2670 files
2671 .into_iter()
2672 .map(|file| async move { self.resolve_file_digest(scope, file).await }),
2673 )
2674 .buffered(MAX_CONCURRENT_FILE_DIGESTS)
2675 .collect()
2676 .await;
2677 Ok(AgentResponse::FileDigestsResolved { resolutions })
2678 }
2679
2680 async fn resolve_file_digest(
2681 &self,
2682 scope: FileDigestScope,
2683 file: FileIdentity,
2684 ) -> FileDigestResolution {
2685 let lock = {
2686 let key = (scope, file.clone());
2687 let mut locks = self.file_digest_locks.lock().unwrap();
2688 locks.retain(|_, lock| lock.strong_count() > 0);
2689 if let Some(lock) = locks.get(&key).and_then(Weak::upgrade) {
2690 lock
2691 } else {
2692 let lock = Arc::new(FileDigestFlight {
2693 lock: tokio::sync::Mutex::new(()),
2694 resolution: Mutex::new(None),
2695 });
2696 locks.insert(key, Arc::downgrade(&lock));
2697 lock
2698 }
2699 };
2700 let _flight = lock.lock.lock().await;
2701 if let Some(resolution) = lock.resolution.lock().unwrap().clone() {
2702 return resolution;
2703 }
2704 if let Some(digest) = self
2705 .file_digests
2706 .lock()
2707 .unwrap()
2708 .get(&(scope, file.path.clone()))
2709 .filter(|recorded| recorded.file == file)
2710 .map(|recorded| recorded.digest.clone())
2711 .filter(|digest| digest.size == file.len && file.still_describes().unwrap_or(false))
2712 {
2713 return FileDigestResolution::Digest(digest);
2714 }
2715 let Ok(_permit) = self.file_digest_permits.acquire().await else {
2716 return FileDigestResolution::Unresolved;
2717 };
2718 #[cfg(test)]
2719 self.file_digest_reads.fetch_add(1, Ordering::Relaxed);
2720 let path = file.path.clone();
2721 let resolved =
2722 tokio::task::spawn_blocking(move || digest_file_validated(scope, &path)).await;
2723 let Ok(Ok(Some(observed))) = resolved else {
2724 return FileDigestResolution::Unresolved;
2725 };
2726 if observed.cache_identity.as_ref() != Some(&file) {
2727 return FileDigestResolution::Unresolved;
2728 }
2729 let resolution = match observed.resolution {
2730 FileDigestResolution::Digest(digest) if digest.size == file.len => {
2731 let _ = self.record_file_digests(
2732 scope,
2733 vec![RecordedFileDigest {
2734 file: file.clone(),
2735 digest: digest.clone(),
2736 }],
2737 );
2738 FileDigestResolution::Digest(digest)
2739 }
2740 FileDigestResolution::EmbeddedTimestampMacro => {
2741 FileDigestResolution::EmbeddedTimestampMacro
2742 }
2743 FileDigestResolution::Digest(_) | FileDigestResolution::Unresolved => {
2744 FileDigestResolution::Unresolved
2745 }
2746 };
2747 *lock.resolution.lock().unwrap() = Some(resolution.clone());
2748 resolution
2749 }
2750
2751 fn record_file_digests(
2757 &self,
2758 scope: FileDigestScope,
2759 entries: Vec<RecordedFileDigest>,
2760 ) -> Result<AgentResponse> {
2761 if entries.len() > MAX_FILE_DIGEST_BATCH {
2762 bail!("too many file-digest records in one request");
2763 }
2764 for entry in &entries {
2765 validate_file_digest_record(entry)?;
2766 }
2767 let mut ledger = self.file_digests.lock().unwrap();
2768 for entry in entries {
2769 if ledger.len() >= MAX_FILE_DIGEST_ENTRIES
2770 && !ledger.contains_key(&(scope, entry.file.path.clone()))
2771 {
2772 break;
2773 }
2774 ledger.insert((scope, entry.file.path.clone()), entry);
2775 }
2776 Ok(AgentResponse::FileDigestsRecorded)
2777 }
2778
2779 pub fn seed_file_digests(&self, entries: Vec<(FileDigestScope, RecordedFileDigest)>) -> usize {
2788 self.seed_file_digests_with(|| entries)
2789 }
2790
2791 pub fn seed_file_digests_with(
2797 &self,
2798 load: impl FnOnce() -> Vec<(FileDigestScope, RecordedFileDigest)>,
2799 ) -> usize {
2800 let mut ledger = self.file_digests.lock().unwrap();
2801 let entries = load();
2802 let mut seeded = 0;
2803 for (scope, entry) in entries {
2804 if ledger.len() >= MAX_FILE_DIGEST_ENTRIES {
2805 break;
2806 }
2807 if validate_file_digest_record(&entry).is_err() {
2808 continue;
2809 }
2810 if let std::collections::btree_map::Entry::Vacant(slot) =
2811 ledger.entry((scope, entry.file.path.clone()))
2812 {
2813 slot.insert(entry);
2814 seeded += 1;
2815 }
2816 }
2817 seeded
2818 }
2819
2820 pub fn file_digests(&self) -> Vec<(FileDigestScope, RecordedFileDigest)> {
2823 self.file_digests
2824 .lock()
2825 .unwrap()
2826 .iter()
2827 .map(|((scope, _), entry)| (*scope, entry.clone()))
2828 .collect()
2829 }
2830
2831 fn find_action_prediction(
2832 &self,
2833 task: &str,
2834 invocation: &CacheDigest,
2835 ) -> Result<AgentResponse> {
2836 validate_task_identity(task)?;
2837 invocation.validate()?;
2838 let (prediction, prefetch) = {
2839 let mut tasks = self.task_actions.lock().unwrap();
2840 let Some(state) = tasks.get_mut(task) else {
2841 return Ok(AgentResponse::ActionPrediction { prediction: None });
2842 };
2843 activate_prediction_adapter(state, invocation)
2844 };
2845 if let Some(prefetch) = prefetch {
2846 self.spawn_prefetch_predictions(prefetch);
2847 }
2848 Ok(AgentResponse::ActionPrediction { prediction })
2849 }
2850
2851 fn record_action_prediction(
2852 &self,
2853 task: &str,
2854 prediction: ActionPrediction,
2855 ) -> Result<AgentResponse> {
2856 validate_task_identity(task)?;
2857 prediction.validate()?;
2858 let mut tasks = self.task_actions.lock().unwrap();
2859 let state = tasks.entry(task.to_string()).or_default();
2860 if !state
2861 .pending_predictions
2862 .contains_key(&prediction.invocation)
2863 && state.pending_predictions.len() >= MAX_TASK_ACTION_PREDICTIONS
2864 {
2865 bail!("this task run contains too many action predictions");
2866 }
2867 let current = state.predictions.get(&prediction.invocation);
2868 if current != Some(&prediction) {
2869 match state.changed_predictions.get(&prediction.invocation) {
2870 Some(baseline) if baseline.as_ref() == Some(&prediction) => {
2872 state.changed_predictions.remove(&prediction.invocation);
2873 }
2874 Some(_) => {}
2876 None => {
2877 state
2878 .changed_predictions
2879 .insert(prediction.invocation.clone(), current.cloned());
2880 }
2881 }
2882 }
2883 state
2884 .predictions
2885 .insert(prediction.invocation.clone(), prediction.clone());
2886 state
2887 .pending_predictions
2888 .insert(prediction.invocation.clone(), prediction);
2889 Ok(AgentResponse::ActionPredictionRecorded)
2890 }
2891
2892 fn executable_identity_key(
2893 &self,
2894 executable: PathBuf,
2895 environment: BTreeMap<String, Option<String>>,
2896 ) -> Result<ExecutableIdentityKey> {
2897 if !environment.keys().all(|name| {
2904 matches!(
2905 name.as_str(),
2906 "COMPILER_PATH"
2907 | "GCC_EXEC_PREFIX"
2908 | "LIBRARY_PATH"
2909 | "MBX_FUSE_LD"
2910 | "MBX_WORKING_DIRECTORY"
2911 | "PATH"
2912 | "RUSTUP_HOME"
2913 | "RUSTUP_TOOLCHAIN"
2914 | "SDKROOT"
2915 | "MACOSX_DEPLOYMENT_TARGET"
2916 | "LIB"
2917 | "UCRTVersion"
2918 | "UniversalCRTSdkDir"
2919 | "VCToolsInstallDir"
2920 | "VCToolsVersion"
2921 | "WindowsSdkDir"
2922 | "WindowsSDKVersion"
2923 )
2924 }) {
2925 bail!("executable identity contains an unsupported environment variable");
2926 }
2927 Ok(ExecutableIdentityKey {
2928 executable,
2929 environment,
2930 })
2931 }
2932
2933 fn find_executable_identity(
2934 &self,
2935 executable: PathBuf,
2936 environment: BTreeMap<String, Option<String>>,
2937 ) -> Result<AgentResponse> {
2938 let key = self.executable_identity_key(executable, environment)?;
2939 let remembered = self
2940 .executable_identities
2941 .lock()
2942 .unwrap()
2943 .get(&key)
2944 .cloned();
2945 let stdout = match remembered {
2946 Some(stdout) => Some(stdout),
2947 None => {
2948 let persisted = self.load_persisted_identity(&key);
2949 if let Some(stdout) = &persisted {
2950 let _ = self.remember_executable_identity(&key, stdout.clone());
2953 }
2954 persisted
2955 }
2956 };
2957 Ok(AgentResponse::ExecutableIdentity { stdout })
2958 }
2959
2960 fn store_executable_identity(
2961 &self,
2962 executable: PathBuf,
2963 environment: BTreeMap<String, Option<String>>,
2964 stdout: Vec<u8>,
2965 pins: Vec<PinnedFile>,
2966 ) -> Result<AgentResponse> {
2967 if stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE {
2968 bail!("executable identity exceeds {MAX_EXECUTABLE_IDENTITY_SIZE} bytes");
2969 }
2970 let key = self.executable_identity_key(executable, environment)?;
2971 self.remember_executable_identity(&key, stdout.clone())?;
2972 if pins.is_empty() {
2977 } else if !pins.iter().all(PinnedFile::holds) {
2978 debug!("executable identity was not persisted: a pinned file changed during the probe");
2979 } else if let Err(error) = self.persist_executable_identity(&key, pins, &stdout) {
2980 debug!("executable identity was not persisted: {error:#}");
2981 }
2982 Ok(AgentResponse::ExecutableIdentity {
2983 stdout: Some(stdout),
2984 })
2985 }
2986
2987 fn remember_executable_identity(
2988 &self,
2989 key: &ExecutableIdentityKey,
2990 stdout: Vec<u8>,
2991 ) -> Result<()> {
2992 let mut identities = self.executable_identities.lock().unwrap();
2993 let is_new = !identities.contains_key(key);
2994 let previous_size = identities.get(key).map_or(0, Vec::len);
2995 if is_new && identities.len() >= MAX_EXECUTABLE_IDENTITIES {
2996 bail!("executable identity cache contains too many entries");
2997 }
2998 let retained_bytes = identities.values().map(Vec::len).sum::<usize>();
2999 if retained_bytes - previous_size + stdout.len() > MAX_EXECUTABLE_IDENTITY_BYTES {
3000 bail!("executable identity cache contains too many bytes");
3001 }
3002 identities.insert(key.clone(), stdout);
3003 Ok(())
3004 }
3005
3006 fn persisted_identity_path(&self, key: &ExecutableIdentityKey) -> Result<PathBuf> {
3007 let selector = serde_json::to_vec(&(&key.executable, &key.environment))?;
3008 let digest = CacheDigest::blake3(&selector);
3009 Ok(self.identity_dir.join(format!("{}.json", digest.hash)))
3010 }
3011
3012 fn load_persisted_identity(&self, key: &ExecutableIdentityKey) -> Option<Vec<u8>> {
3015 let path = self.persisted_identity_path(key).ok()?;
3016 let bytes = fs::read(path).ok()?;
3017 let persisted: PersistedExecutableIdentity = serde_json::from_slice(&bytes).ok()?;
3018 if persisted.version != PERSISTED_EXECUTABLE_IDENTITY_VERSION
3019 || persisted.executable != key.executable
3020 || persisted.environment != key.environment
3021 || persisted.pins.is_empty()
3022 || persisted.stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE
3023 {
3024 return None;
3025 }
3026 persisted
3027 .pins
3028 .iter()
3029 .all(PinnedFile::holds)
3030 .then(|| persisted.stdout.into_bytes())
3031 }
3032
3033 fn persist_executable_identity(
3034 &self,
3035 key: &ExecutableIdentityKey,
3036 pins: Vec<PinnedFile>,
3037 stdout: &[u8],
3038 ) -> Result<()> {
3039 let Ok(stdout) = std::str::from_utf8(stdout) else {
3041 return Ok(());
3042 };
3043 let persisted = PersistedExecutableIdentity {
3044 version: PERSISTED_EXECUTABLE_IDENTITY_VERSION,
3045 executable: key.executable.clone(),
3046 environment: key.environment.clone(),
3047 pins,
3048 stdout: stdout.to_owned(),
3049 };
3050 let path = self.persisted_identity_path(key)?;
3051 fs::create_dir_all(self.identity_dir.as_path())?;
3052 let mut temporary = tempfile::NamedTempFile::new_in(self.identity_dir.as_path())?;
3054 std::io::Write::write_all(temporary.as_file_mut(), &serde_json::to_vec(&persisted)?)?;
3055 temporary.persist(path).map_err(|error| error.error)?;
3056 Ok(())
3057 }
3058
3059 pub async fn handle_connection<S>(&self, stream: S) -> Result<()>
3061 where
3062 S: AsyncRead + AsyncWrite + Unpin,
3063 {
3064 let (reader, mut writer) = tokio::io::split(stream);
3065 let mut reader = BufReader::new(reader);
3066 let hello = read_request(&mut reader)
3067 .await?
3068 .ok_or_else(|| eyre::eyre!("connection closed before the agent handshake"))?;
3069 let request: AgentRequest = serde_json::from_str(&hello)?;
3070 match request {
3071 AgentRequest::Hello {
3072 protocol,
3073 client_version,
3074 } if protocol == AGENT_PROTOCOL_VERSION && client_version == self.version.as_ref() => {}
3075 AgentRequest::Hello { protocol, .. } if protocol != AGENT_PROTOCOL_VERSION => {
3076 send_response(
3077 &mut writer,
3078 &AgentResponse::Error {
3079 message: format!(
3080 "unsupported agent protocol {protocol}; expected {AGENT_PROTOCOL_VERSION}"
3081 ),
3082 },
3083 )
3084 .await?;
3085 return Ok(());
3086 }
3087 AgentRequest::Hello { client_version, .. } => {
3088 send_response(
3089 &mut writer,
3090 &AgentResponse::Error {
3091 message: format!(
3092 "cache client {client_version} does not match agent {}",
3093 self.version
3094 ),
3095 },
3096 )
3097 .await?;
3098 return Ok(());
3099 }
3100 _ => bail!("the first agent request must be hello"),
3101 }
3102 send_response(
3103 &mut writer,
3104 &AgentResponse::Hello {
3105 protocol: AGENT_PROTOCOL_VERSION,
3106 agent_version: self.version.to_string(),
3107 },
3108 )
3109 .await?;
3110
3111 let mut connection = ConnectionUploads::default();
3115 let served = async {
3116 while let Some(line) = read_request(&mut reader).await? {
3117 let response = match serde_json::from_str(&line) {
3118 Ok(request) => self.respond_on(request, &mut connection).await,
3119 Err(error) => AgentResponse::Error {
3120 message: format!("invalid agent request: {error}"),
3121 },
3122 };
3123 send_response(&mut writer, &response).await?;
3124 }
3125 Ok(())
3126 }
3127 .await;
3128 self.release_bypass(&mut connection);
3130 served
3131 }
3132}
3133
3134async fn read_request<R>(reader: &mut R) -> Result<Option<String>>
3140where
3141 R: AsyncBufRead + Unpin,
3142{
3143 let mut line = Vec::new();
3144 loop {
3145 let available = reader.fill_buf().await?;
3146 if available.is_empty() {
3147 break;
3148 }
3149 let (consumed, complete) = match available.iter().position(|byte| *byte == b'\n') {
3150 Some(index) => (index, true),
3151 None => (available.len(), false),
3152 };
3153 if line.len() + consumed > MAX_REQUEST_BYTES {
3154 bail!("agent request exceeded {MAX_REQUEST_BYTES} bytes");
3155 }
3156 line.extend_from_slice(&available[..consumed]);
3157 reader.consume(consumed + usize::from(complete));
3159 if complete {
3160 return Ok(Some(String::from_utf8(line)?));
3161 }
3162 }
3163 if line.is_empty() {
3164 Ok(None)
3165 } else {
3166 Ok(Some(String::from_utf8(line)?))
3167 }
3168}
3169
3170async fn send_response(
3171 writer: &mut (impl AsyncWrite + Unpin),
3172 response: &AgentResponse,
3173) -> Result<()> {
3174 let mut encoded = serde_json::to_vec(response)?;
3175 encoded.push(b'\n');
3176 writer.write_all(&encoded).await?;
3177 writer.flush().await?;
3178 Ok(())
3179}
3180
3181#[cfg(test)]
3182#[path = "agent_tests.rs"]
3183mod tests;