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