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