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