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