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::{info, warn};
11use std::collections::{BTreeMap, BTreeSet};
12use std::fs;
13use std::path::{Path, PathBuf};
14use std::sync::atomic::{AtomicU64, Ordering};
15use std::sync::{Arc, Mutex, Weak};
16use std::time::{Duration, Instant, SystemTime};
17
18const ACTION_PROMISE_WAIT: Duration = Duration::from_secs(60 * 60);
22use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
23
24mod file_digest;
25mod manifest;
26mod prefetch;
27mod stats;
28mod wire;
29
30#[cfg(test)]
31pub(crate) use prefetch::select_prefetch_actions;
32
33pub use file_digest::{
34 FileDigestCache, FileDigestScope, FileIdentity, FileSnapshot, NoFileDigestCache,
35 RecordedFileDigest,
36};
37pub use manifest::{is_task_identity, task_manifest_actions};
38use manifest::{
39 merge_remote_task_manifest, merge_task_manifests, task_manifest_dir, validate_task_identity,
40 validate_task_manifest,
41};
42pub use stats::{AgentStats, CompilerStats};
43use wire::MAX_REQUEST_BYTES;
44pub use wire::{
45 AGENT_PROTOCOL_VERSION, ActionDiagnostic, AgentEvent, AgentEventObserver, AgentRequest,
46 AgentResponse, RestoreStats,
47};
48
49const MAX_EXECUTABLE_IDENTITIES: usize = 64;
50const MAX_EXECUTABLE_IDENTITY_SIZE: usize = 64 * 1024;
51const MAX_EXECUTABLE_IDENTITY_BYTES: usize = 256 * 1024;
52const TASK_ACTION_MANIFEST_VERSION: u8 = 1;
53const MAX_TASK_ACTION_PREDICTIONS: usize = 16 * 1024;
54const MAX_WARNING_BYTES: usize = 4 * 1024;
56const MAX_WARNINGS: usize = 128;
62const ACTION_DIAGNOSTIC_PREFIX: &str = "@mbx-action-diagnostic\t";
63const MAX_FILE_DIGEST_BATCH: usize = 16 * 1024;
65const MAX_FILE_DIGEST_ENTRIES: usize = 1024 * 1024;
70const MAX_REMOTE_TRANSFERS: usize = 64;
71const MAX_PREFETCH_TRANSFERS: usize = 48;
72const MAX_CONCURRENT_BLOB_PACKS: usize = 8;
78const MAX_PREFETCH_ACTIONS: usize = 1024;
89const MAX_PREFETCH_ACTION_WAVE: usize = 32;
94const MAX_PREFETCH_BATCH_LOOKUPS: usize = 1;
100const PREFETCH_ACTION_BATCH_DELAY: Duration = Duration::from_millis(5);
101const MAX_PREFETCH_DIRECTORY_OBJECTS: usize = 100_000;
102const MAX_PREFETCH_OBJECTS_PER_WAVE: usize = 100_000;
103const DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES: u64 = 5 * 1024 * 1024 * 1024;
104
105pub type TaskFallbacks = Arc<dyn Fn() -> Vec<String> + Send + Sync>;
108
109const NEWEST_MANIFEST_CANDIDATES: usize = 8;
121
122pub struct AgentRemoteCache {
124 pub client: RemoteCacheClient,
126 pub mode: RemoteCacheMode,
128 pub staging_dir: PathBuf,
130}
131
132#[derive(Default)]
133struct AtomicAgentStats {
134 lookups: AtomicU64,
135 unconsulted: AtomicU64,
136 hits: AtomicU64,
137 stores: AtomicU64,
138 stored_bytes: AtomicU64,
139 verifications: AtomicU64,
140 divergences: AtomicU64,
141 downloaded_bytes: AtomicU64,
142 uploaded_bytes: AtomicU64,
143 background_uploads: AtomicU64,
144 background_upload_failures: AtomicU64,
145 remote_blob_pack_uploads: AtomicU64,
146 remote_blob_pack_upload_blobs: AtomicU64,
147 upload_drain_duration_ns: AtomicU64,
148 prefetched_actions: AtomicU64,
149 predictions_loaded: AtomicU64,
150 remote_failures: AtomicU64,
151 remote_manifest_lookups: AtomicU64,
152 remote_manifest_lookup_duration_ns: AtomicU64,
153 remote_action_lookups: AtomicU64,
154 remote_action_lookup_duration_ns: AtomicU64,
155 remote_blob_requests: AtomicU64,
156 remote_blob_pack_requests: AtomicU64,
157 remote_blob_pack_blobs: AtomicU64,
158 remote_blob_transfer_duration_ns: AtomicU64,
159 local_cas_write_duration_ns: AtomicU64,
160 prefetch_runs: AtomicU64,
161 prefetch_duration_ns: AtomicU64,
162 materialization_duration_ns: AtomicU64,
163 bypasses: Mutex<BTreeMap<String, u64>>,
164 avoided_compiler_duration_ns: AtomicU64,
165 compiler: Mutex<BTreeMap<String, CompilerStats>>,
166 slow_compilations: Mutex<BTreeMap<String, u64>>,
167 restored_output_files: AtomicU64,
168 restored_output_bytes: AtomicU64,
169 reflinked_output_files: AtomicU64,
170 reflinked_output_bytes: AtomicU64,
171 copied_output_files: AtomicU64,
172 copied_output_bytes: AtomicU64,
173 reused_output_files: AtomicU64,
174 reused_output_bytes: AtomicU64,
175}
176
177struct AtomicDurationTimer<'a> {
178 started: Instant,
179 target: &'a AtomicU64,
180}
181
182impl<'a> AtomicDurationTimer<'a> {
183 fn start(target: &'a AtomicU64) -> Self {
184 Self {
185 started: Instant::now(),
186 target,
187 }
188 }
189}
190
191impl Drop for AtomicDurationTimer<'_> {
192 fn drop(&mut self) {
193 atomic_saturating_add(self.target, duration_ns(self.started));
194 }
195}
196
197fn duration_ns(started: Instant) -> u64 {
198 started.elapsed().as_nanos().try_into().unwrap_or(u64::MAX)
199}
200
201fn validate_crate_name(crate_name: Option<&str>) -> Result<()> {
202 if let Some(crate_name) = crate_name
203 && (crate_name.len() > 256 || crate_name.contains(['\0', '\n', '\r']))
204 {
205 bail!("invalid compiler crate name");
206 }
207 Ok(())
208}
209
210#[derive(serde::Deserialize)]
211struct ActionDiagnosticEnvelope {
212 outcome: String,
213 crate_name: Option<String>,
214 diagnostic: ActionDiagnostic,
215}
216
217fn parse_action_diagnostic(
218 message: &str,
219) -> Option<Result<(String, Option<String>, ActionDiagnostic)>> {
220 let payload = message.strip_prefix(ACTION_DIAGNOSTIC_PREFIX)?;
221 Some((|| {
222 let envelope: ActionDiagnosticEnvelope = serde_json::from_str(payload)?;
223 if !matches!(envelope.outcome.as_str(), "hit" | "miss") {
224 bail!("invalid action diagnostic outcome");
225 }
226 validate_crate_name(envelope.crate_name.as_deref())?;
227 Ok((envelope.outcome, envelope.crate_name, envelope.diagnostic))
228 })())
229}
230
231fn atomic_saturating_add(target: &AtomicU64, value: u64) {
232 let _ = target.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
233 Some(current.saturating_add(value))
234 });
235}
236
237fn queue_prefetch_digest(
238 verified: &BTreeMap<CacheDigest, PathBuf>,
239 pending: &mut BTreeMap<CacheDigest, ()>,
240 digest: CacheDigest,
241) {
242 if verified.contains_key(&digest) || pending.contains_key(&digest) {
243 return;
244 }
245 pending.insert(digest, ());
246}
247
248fn queue_prefetch_directory(
249 seen: &BTreeMap<CacheDigest, ()>,
250 pending: &mut BTreeMap<CacheDigest, ()>,
251 digest: CacheDigest,
252 limit: usize,
253) -> bool {
254 if seen.contains_key(&digest) || pending.contains_key(&digest) {
255 return true;
256 }
257 if seen.len().saturating_add(pending.len()) >= limit {
258 return false;
259 }
260 pending.insert(digest, ());
261 true
262}
263
264pub type TaskLoader = dyn for<'a> Fn(&'a CacheAgent, &'a str) -> BoxFuture<'a, ()> + Send + Sync;
266
267#[derive(Clone)]
273pub struct CacheAgent {
274 cas: LocalCas,
275 actions: LocalActionCache,
276 verified_blobs: Arc<Mutex<BTreeMap<CacheDigest, VerifiedBlob>>>,
277 version: Arc<str>,
278 write_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
279 action_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
280 stats: Arc<AtomicAgentStats>,
281 observer: Option<Arc<dyn AgentEventObserver>>,
282 observer_emission: Arc<Mutex<()>>,
283 task_loader: Option<Arc<TaskLoader>>,
286 executable_identities: Arc<Mutex<BTreeMap<ExecutableIdentityKey, Vec<u8>>>>,
287 manifest_dir: Arc<PathBuf>,
288 task_actions: Arc<Mutex<BTreeMap<String, TaskActionState>>>,
289 task_fallbacks: Arc<Mutex<BTreeMap<String, TaskFallbacks>>>,
291 next_task_run: Arc<AtomicU64>,
292 manifest_write_lock: Arc<Mutex<()>>,
293 remote: Option<Arc<RemoteCacheClient>>,
294 remote_mode: RemoteCacheMode,
295 remote_staging_dir: Arc<PathBuf>,
296 remote_download_limit: u64,
297 remote_download_bytes: Arc<AtomicU64>,
298 pending_remote_actions: Arc<Mutex<BTreeMap<CacheDigest, RemoteActionResult>>>,
299 remote_transfers: Arc<tokio::sync::Semaphore>,
300 prefetch_transfers: Arc<tokio::sync::Semaphore>,
301 prefetch_tasks: Arc<Mutex<Vec<tokio::task::JoinHandle<()>>>>,
302 warnings: Arc<Mutex<BTreeSet<String>>>,
305 file_digests: Arc<Mutex<BTreeMap<(FileDigestScope, PathBuf), RecordedFileDigest>>>,
308 uploads: Option<UploadQueue>,
310}
311
312fn validate_file_digest_record(entry: &RecordedFileDigest) -> Result<()> {
315 if !entry.file.path.is_absolute() {
316 bail!("file-digest records need absolute paths");
317 }
318 entry.digest.validate()?;
319 if entry.file.len != entry.digest.size {
320 bail!("file-digest record length does not match its digest");
321 }
322 Ok(())
323}
324
325struct AgentUploadSink {
327 stats: Arc<AtomicAgentStats>,
328}
329
330impl UploadSink for AgentUploadSink {
331 fn record_blob_uploaded(&self, bytes: u64) {
332 self.stats
333 .background_uploads
334 .fetch_add(1, Ordering::Relaxed);
335 self.stats
336 .uploaded_bytes
337 .fetch_add(bytes, Ordering::Relaxed);
338 }
339
340 fn record_action_uploaded(&self) {
341 self.stats
342 .background_uploads
343 .fetch_add(1, Ordering::Relaxed);
344 }
345
346 fn record_blob_pack_uploaded(&self, blobs: u64) {
347 self.stats
348 .remote_blob_pack_uploads
349 .fetch_add(1, Ordering::Relaxed);
350 self.stats
351 .remote_blob_pack_upload_blobs
352 .fetch_add(blobs, Ordering::Relaxed);
353 }
354
355 fn record_upload_failure(&self) {
356 self.stats
357 .background_upload_failures
358 .fetch_add(1, Ordering::Relaxed);
359 self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
360 }
361}
362
363#[derive(Debug, Clone)]
369struct VerifiedBlob {
370 path: PathBuf,
371 len: u64,
372 modified: SystemTime,
373}
374
375impl VerifiedBlob {
376 fn describe(path: &Path) -> Option<Self> {
379 let metadata = std::fs::metadata(path).ok()?;
380 Some(Self {
381 path: path.to_path_buf(),
382 len: metadata.len(),
383 modified: metadata.modified().ok()?,
384 })
385 }
386
387 fn is_unchanged(&self) -> bool {
389 let Ok(metadata) = std::fs::metadata(&self.path) else {
390 return false;
391 };
392 metadata.len() == self.len && metadata.modified().is_ok_and(|now| now == self.modified)
393 }
394}
395
396#[derive(Debug, Clone, Default)]
397struct TaskActionState {
398 manifest: String,
399 baseline_loaded: bool,
400 predictions: BTreeMap<CacheDigest, ActionPrediction>,
401 pending_predictions: BTreeMap<CacheDigest, ActionPrediction>,
402 prefetched_adapters: BTreeSet<String>,
403 remote_etag: Option<String>,
404}
405
406fn activate_prediction_adapter(
408 state: &mut TaskActionState,
409 invocation: &CacheDigest,
410) -> (Option<ActionPrediction>, Option<Vec<ActionPrediction>>) {
411 let prediction = state.predictions.get(invocation).cloned();
412 let prefetch = prediction.as_ref().and_then(|prediction| {
413 if !state.prefetched_adapters.insert(prediction.adapter.clone()) {
414 return None;
415 }
416 Some(
417 state
418 .predictions
419 .values()
420 .filter(|candidate| candidate.adapter == prediction.adapter)
421 .cloned()
422 .collect(),
423 )
424 });
425 (prediction, prefetch)
426}
427
428struct PrefetchedAction {
429 adapter: String,
430 result: RemoteActionResult,
431}
432
433struct RemoteDownloadReservation {
434 counter: Arc<AtomicU64>,
435 reserved: u64,
436 committed: bool,
437}
438
439impl RemoteDownloadReservation {
440 fn bytes(&self) -> u64 {
441 self.reserved
442 }
443
444 fn commit(mut self, bytes: u64) {
445 debug_assert!(bytes <= self.reserved);
446 self.counter
447 .fetch_sub(self.reserved.saturating_sub(bytes), Ordering::AcqRel);
448 self.committed = true;
449 }
450}
451
452impl Drop for RemoteDownloadReservation {
453 fn drop(&mut self) {
454 if !self.committed {
455 self.counter.fetch_sub(self.reserved, Ordering::AcqRel);
456 }
457 }
458}
459
460#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
461struct ExecutableIdentityKey {
462 executable: PathBuf,
463 environment: BTreeMap<String, Option<String>>,
464}
465
466impl CacheAgent {
467 pub fn new(cache_dir: impl Into<PathBuf>, version: impl Into<Arc<str>>) -> Self {
469 Self::build(cache_dir.into(), version.into(), None, 0)
470 }
471
472 pub fn new_remote(
474 cache_dir: impl Into<PathBuf>,
475 version: impl Into<Arc<str>>,
476 remote: AgentRemoteCache,
477 ) -> Self {
478 Self::build(
479 cache_dir.into(),
480 version.into(),
481 Some(remote),
482 DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES,
483 )
484 }
485
486 pub fn new_remote_with_download_limit(
488 cache_dir: impl Into<PathBuf>,
489 version: impl Into<Arc<str>>,
490 remote: AgentRemoteCache,
491 max_remote_download_bytes: u64,
492 ) -> Self {
493 Self::build(
494 cache_dir.into(),
495 version.into(),
496 Some(remote),
497 max_remote_download_bytes,
498 )
499 }
500
501 fn build(
502 cache_dir: PathBuf,
503 version: Arc<str>,
504 remote: Option<AgentRemoteCache>,
505 remote_download_limit: u64,
506 ) -> Self {
507 let remote_mode = remote
508 .as_ref()
509 .map_or(RemoteCacheMode::ReadOnly, |remote| remote.mode);
510 let remote_staging_dir = remote.as_ref().map_or_else(
511 || cache_dir.join("remote"),
512 |remote| remote.staging_dir.clone(),
513 );
514 let remote = remote.map(|remote| Arc::new(remote.client));
515 let stats = Arc::new(AtomicAgentStats::default());
516 let remote_transfers = Arc::new(tokio::sync::Semaphore::new(MAX_REMOTE_TRANSFERS));
517 let uploads = remote
520 .clone()
521 .filter(|_| remote_mode.writes())
522 .map(|client| {
523 UploadQueue::new(
524 client,
525 Arc::new(AgentUploadSink {
526 stats: stats.clone(),
527 }),
528 remote_transfers.clone(),
529 )
530 });
531 Self {
532 cas: LocalCas::new(cache_dir.clone()),
533 actions: LocalActionCache::new(cache_dir.clone()),
534 verified_blobs: Arc::new(Mutex::new(BTreeMap::new())),
535 version,
536 write_locks: Arc::new(Mutex::new(BTreeMap::new())),
537 action_locks: Arc::new(Mutex::new(BTreeMap::new())),
538 stats,
539 observer: None,
540 observer_emission: Arc::new(Mutex::new(())),
541 task_loader: None,
542 executable_identities: Arc::new(Mutex::new(BTreeMap::new())),
543 manifest_dir: Arc::new(task_manifest_dir(&cache_dir)),
544 task_actions: Arc::new(Mutex::new(BTreeMap::new())),
545 task_fallbacks: Arc::new(Mutex::new(BTreeMap::new())),
546 next_task_run: Arc::new(AtomicU64::new(0)),
547 manifest_write_lock: Arc::new(Mutex::new(())),
548 remote,
549 remote_mode,
550 remote_staging_dir: Arc::new(remote_staging_dir),
551 remote_download_limit,
552 remote_download_bytes: Arc::new(AtomicU64::new(0)),
553 pending_remote_actions: Arc::new(Mutex::new(BTreeMap::new())),
554 remote_transfers,
555 prefetch_transfers: Arc::new(tokio::sync::Semaphore::new(MAX_PREFETCH_TRANSFERS)),
556 prefetch_tasks: Arc::new(Mutex::new(Vec::new())),
557 warnings: Arc::new(Mutex::new(BTreeSet::new())),
558 file_digests: Arc::new(Mutex::new(BTreeMap::new())),
559 uploads,
560 }
561 }
562
563 #[must_use]
565 pub fn with_observer(mut self, observer: Arc<dyn AgentEventObserver>) -> Self {
566 self.observer = Some(observer);
567 self
568 }
569
570 pub fn with_task_loader(mut self, loader: Arc<TaskLoader>) -> Self {
577 self.task_loader = Some(loader);
578 self
579 }
580
581 async fn ensure_task_loaded(&self, task: &str) {
584 let Some(loader) = &self.task_loader else {
585 return;
586 };
587 if self.task_actions.lock().unwrap().contains_key(task) {
588 return;
589 }
590 loader(self, task).await;
591 }
592
593 fn emit(&self, event: impl FnOnce() -> AgentEvent) {
594 if let Some(observer) = &self.observer {
595 let _emission = self.observer_emission.lock().unwrap();
596 observer.event(event());
597 }
598 }
599
600 fn emit_action(&self, diagnostic: Option<AgentEvent>, action: AgentEvent) {
601 if let Some(observer) = &self.observer {
602 let _emission = self.observer_emission.lock().unwrap();
603 if let Some(diagnostic) = diagnostic {
604 observer.event(diagnostic);
605 }
606 observer.event(action);
607 }
608 }
609
610 fn reserve_remote_download(&self, bytes: u64) -> Result<RemoteDownloadReservation> {
611 let mut current = self.remote_download_bytes.load(Ordering::Acquire);
612 loop {
613 let next = current
614 .checked_add(bytes)
615 .ok_or_else(|| eyre::eyre!("remote cache download budget overflowed"))?;
616 if next > self.remote_download_limit {
617 bail!(
618 "remote cache download budget exceeded: {} bytes requested with {} of {} bytes already used",
619 bytes,
620 current,
621 self.remote_download_limit
622 );
623 }
624 match self.remote_download_bytes.compare_exchange_weak(
625 current,
626 next,
627 Ordering::AcqRel,
628 Ordering::Acquire,
629 ) {
630 Ok(_) => {
631 return Ok(RemoteDownloadReservation {
632 counter: self.remote_download_bytes.clone(),
633 reserved: bytes,
634 committed: false,
635 });
636 }
637 Err(observed) => current = observed,
638 }
639 }
640 }
641
642 fn reserve_remote_download_up_to(&self, requested: u64) -> Result<RemoteDownloadReservation> {
643 let mut current = self.remote_download_bytes.load(Ordering::Acquire);
644 loop {
645 let reserved = requested.min(self.remote_download_limit.saturating_sub(current));
646 let next = current + reserved;
647 match self.remote_download_bytes.compare_exchange_weak(
648 current,
649 next,
650 Ordering::AcqRel,
651 Ordering::Acquire,
652 ) {
653 Ok(_) => {
654 return Ok(RemoteDownloadReservation {
655 counter: self.remote_download_bytes.clone(),
656 reserved,
657 committed: false,
658 });
659 }
660 Err(observed) => current = observed,
661 }
662 }
663 }
664
665 pub async fn begin_task(&self, task: &str) -> Result<String> {
667 self.begin_task_with_remote_errors(task, false, None, true)
668 .await
669 }
670
671 pub async fn begin_task_on_prediction(&self, task: &str) -> Result<String> {
674 self.begin_task_with_remote_errors(task, false, None, false)
675 .await
676 }
677
678 pub async fn begin_session_task(&self, task: &str) -> Result<()> {
684 self.begin_task_with_remote_errors(task, false, Some(task.to_string()), false)
685 .await?;
686 Ok(())
687 }
688
689 pub fn register_task_fallbacks(
700 &self,
701 task: &str,
702 fallbacks: impl Fn() -> Vec<String> + Send + Sync + 'static,
703 ) -> Result<()> {
704 validate_task_identity(task)?;
705 self.task_fallbacks
706 .lock()
707 .unwrap()
708 .insert(task.to_string(), Arc::new(fallbacks));
709 Ok(())
710 }
711
712 pub async fn prefetch_task(&self, task: &str) -> Result<String> {
714 let run = self
715 .begin_task_with_remote_errors(task, true, None, true)
716 .await?;
717 self.wait_for_prefetches().await;
718 Ok(run)
719 }
720
721 async fn begin_task_with_remote_errors(
722 &self,
723 task: &str,
724 strict: bool,
725 run: Option<String>,
726 eager_prefetch: bool,
727 ) -> Result<String> {
728 validate_task_identity(task)?;
729 let early_manifest = {
735 let _write_guard = self.manifest_write_lock.lock().unwrap();
736 let _file_guard = self.lock_task_manifest(task)?;
737 self.load_task_manifest(task)?
738 };
739 let early_actions: BTreeSet<_> = early_manifest
740 .iter()
741 .flat_map(|manifest| manifest.predictions.iter())
742 .map(|prediction| prediction.action.clone())
743 .collect();
744 if eager_prefetch {
745 self.spawn_prefetch_predictions(
746 early_manifest
747 .as_ref()
748 .map(|manifest| manifest.predictions.clone())
749 .unwrap_or_default(),
750 );
751 }
752 let (remote_manifest, mut remote_etag) = if self.remote_mode.reads() {
753 match self.get_remote_task_manifest(task).await {
754 Ok(Some((manifest, etag))) => (Some(manifest), Some(etag)),
755 Ok(None) => (None, None),
756 Err(error) => {
757 if strict {
758 return Err(error).wrap_err_with(|| {
759 format!("remote task action manifest lookup failed for {task}")
760 });
761 }
762 self.note_remote_failure();
763 warn!("remote task action manifest lookup failed for {task}: {error}");
764 (None, None)
765 }
766 }
767 } else {
768 (None, None)
769 };
770 let manifest = if self.remote.is_some() && self.remote_mode.reads() {
774 let _write_guard = self.manifest_write_lock.lock().unwrap();
775 let _file_guard = self.lock_task_manifest(task)?;
776 let local_manifest = self.load_task_manifest(task)?;
777 let manifest = match (remote_manifest, local_manifest) {
778 (Some(remote), Some(local)) => {
779 let (manifest, merged) = merge_remote_task_manifest(task, remote, local);
780 if !merged {
781 remote_etag = None;
782 }
783 Some(manifest)
784 }
785 (Some(remote), None) => Some(remote),
786 (None, local) => local,
787 };
788 if let Some(manifest) = &manifest {
789 self.persist_task_manifest(manifest)?;
790 }
791 manifest
792 } else {
793 early_manifest
794 };
795 let (manifest, remote_etag) = match manifest {
796 Some(manifest) => (Some(manifest), remote_etag),
797 None => (self.inherit_task_manifest(task, strict).await?, None),
800 };
801 let mut state = if let Some(manifest) = manifest {
802 TaskActionState {
803 manifest: task.to_string(),
804 baseline_loaded: true,
805 predictions: manifest
806 .predictions
807 .into_iter()
808 .map(|prediction| (prediction.invocation.clone(), prediction))
809 .collect(),
810 pending_predictions: BTreeMap::new(),
811 prefetched_adapters: BTreeSet::new(),
812 remote_etag,
813 }
814 } else {
815 TaskActionState {
816 manifest: task.to_string(),
817 baseline_loaded: true,
818 remote_etag,
819 ..TaskActionState::default()
820 }
821 };
822 let run = run.unwrap_or_else(|| {
823 let sequence = self.next_task_run.fetch_add(1, Ordering::Relaxed);
824 CacheDigest::blake3(format!("{task}\0{}\0{sequence}", std::process::id()).as_bytes())
825 .hash
826 });
827 self.stats.predictions_loaded.fetch_max(
831 state.predictions.len().try_into().unwrap_or(u64::MAX),
832 Ordering::Relaxed,
833 );
834 if eager_prefetch {
835 state.prefetched_adapters.extend(
836 state
837 .predictions
838 .values()
839 .map(|prediction| prediction.adapter.clone()),
840 );
841 }
842 let predictions = state
846 .predictions
847 .values()
848 .filter(|prediction| !early_actions.contains(&prediction.action))
849 .cloned()
850 .collect();
851 self.task_actions.lock().unwrap().insert(run.clone(), state);
852 if eager_prefetch {
853 self.spawn_prefetch_predictions(predictions);
854 }
855 Ok(run)
856 }
857
858 async fn inherit_task_manifest(
866 &self,
867 task: &str,
868 strict: bool,
869 ) -> Result<Option<TaskActionManifest>> {
870 let fallbacks = self.task_fallbacks.lock().unwrap().get(task).cloned();
871 let Some(fallbacks) = fallbacks else {
872 return Ok(None);
873 };
874 let named = tokio::task::spawn_blocking(move || fallbacks()).await?;
875 let mut tried = BTreeSet::from([task.to_string()]);
876 let candidates = named.into_iter().map(|identity| (identity, true)).chain(
877 self.newest_task_identities()
878 .into_iter()
879 .map(|identity| (identity, false)),
880 );
881 for (identity, named) in candidates {
882 if !tried.insert(identity.clone()) || validate_task_identity(&identity).is_err() {
883 continue;
884 }
885 let mut found = self.load_task_manifest(&identity)?;
886 if found.is_none() && named && self.remote_mode.reads() {
887 found = match self.get_remote_task_manifest(&identity).await {
888 Ok(manifest) => manifest.map(|(manifest, _)| manifest),
889 Err(error) => {
890 if strict {
891 return Err(error).wrap_err_with(|| {
892 format!("remote task action manifest lookup failed for {identity}")
893 });
894 }
895 self.note_remote_failure();
896 warn!("remote task action manifest lookup failed for {identity}: {error}");
897 None
898 }
899 };
900 }
901 let Some(mut manifest) = found else {
902 continue;
903 };
904 if manifest.predictions.is_empty()
905 || (!named && manifest.predictions.len() > MAX_TASK_ACTION_PREDICTIONS / 2)
906 {
907 continue;
908 }
909 info!(
910 "no predictions were recorded for {task}; inheriting {} from {identity}",
911 manifest.predictions.len()
912 );
913 manifest.task = task.to_string();
914 validate_task_manifest(&manifest, task)?;
915 let _write_guard = self.manifest_write_lock.lock().unwrap();
916 let _file_guard = self.lock_task_manifest(task)?;
917 if let Some(recorded) = self.load_task_manifest(task)? {
921 return Ok(Some(recorded));
922 }
923 self.persist_task_manifest(&manifest)?;
924 return Ok(Some(manifest));
925 }
926 Ok(None)
927 }
928
929 fn newest_task_identities(&self) -> Vec<String> {
931 let Ok(entries) = fs::read_dir(self.manifest_dir.as_path()) else {
932 return Vec::new();
933 };
934 let mut manifests: Vec<(std::time::SystemTime, String)> = entries
935 .filter_map(|entry| {
936 let entry = entry.ok()?;
937 let name = entry.file_name();
938 let identity = name.to_str()?.strip_suffix(".json")?.to_string();
939 validate_task_identity(&identity).ok()?;
940 let modified = entry.metadata().ok()?.modified().ok()?;
941 Some((modified, identity))
942 })
943 .collect();
944 manifests.sort_by(|left, right| right.cmp(left));
945 manifests
946 .into_iter()
947 .take(NEWEST_MANIFEST_CANDIDATES)
948 .map(|(_, identity)| identity)
949 .collect()
950 }
951
952 pub async fn cancel_prefetches(&self) {
954 let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
955 for task in &tasks {
956 task.abort();
957 }
958 for task in tasks {
959 if let Err(error) = task.await
960 && !error.is_cancelled()
961 {
962 warn!("remote action prefetch task failed: {error}");
963 }
964 }
965 }
966
967 pub async fn wait_for_uploads(&self) {
974 let Some(uploads) = &self.uploads else {
975 return;
976 };
977 let _timer = AtomicDurationTimer::start(&self.stats.upload_drain_duration_ns);
978 uploads.drain().await;
979 }
980
981 async fn wait_for_prefetches(&self) {
982 let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
983 for task in tasks {
984 if let Err(error) = task.await {
985 warn!("remote action prefetch task failed: {error}");
986 }
987 }
988 }
989
990 pub async fn commit_task(&self, run: &str) -> Result<()> {
992 self.commit_task_actions(run).await.map(|_| ())
993 }
994
995 pub async fn commit_task_actions(&self, run: &str) -> Result<Vec<ActionPrediction>> {
1001 validate_task_identity(run)?;
1002 let state = {
1003 let mut runs = self.task_actions.lock().unwrap();
1004 let state = runs
1005 .get(run)
1006 .ok_or_else(|| eyre::eyre!("task action manifest baseline was not loaded"))?;
1007 if !state.baseline_loaded {
1008 bail!("task action manifest baseline was not loaded");
1009 }
1010 if state.pending_predictions.is_empty()
1016 && (self.remote.is_none() || !self.remote_mode.writes())
1017 {
1018 runs.remove(run);
1019 return Ok(Vec::new());
1020 }
1021 state.clone()
1022 };
1023 let task = state.manifest;
1024 validate_task_identity(&task)?;
1025 let completed = state
1026 .pending_predictions
1027 .values()
1028 .cloned()
1029 .collect::<Vec<_>>();
1030 let (manifest, introduced) = {
1031 let _write_guard = self.manifest_write_lock.lock().unwrap();
1032 let _file_guard = self.lock_task_manifest(&task)?;
1033 let mut predictions = self
1034 .load_task_manifest(&task)?
1035 .map(|manifest| {
1036 manifest
1037 .predictions
1038 .into_iter()
1039 .map(|prediction| (prediction.invocation.clone(), prediction))
1040 .collect::<BTreeMap<_, _>>()
1041 })
1042 .unwrap_or_default();
1043 let introduced: BTreeSet<CacheDigest> = state
1048 .pending_predictions
1049 .keys()
1050 .filter(|invocation| !predictions.contains_key(*invocation))
1051 .cloned()
1052 .collect();
1053 predictions.extend(state.pending_predictions);
1058 let manifest = TaskActionManifest {
1059 version: TASK_ACTION_MANIFEST_VERSION,
1060 task: task.clone(),
1061 predictions: predictions.into_values().collect(),
1062 };
1063 validate_task_manifest(&manifest, &task)?;
1064 self.persist_task_manifest(&manifest)?;
1065 (manifest, introduced)
1066 };
1067 self.task_actions.lock().unwrap().remove(run);
1068 if self.remote_mode.writes() {
1069 let mut manifest = manifest;
1073 if let Some(uploads) = &self.uploads {
1074 let actions: Vec<CacheDigest> = manifest
1075 .predictions
1076 .iter()
1077 .map(|prediction| prediction.action.clone())
1078 .collect();
1079 let unpublished = uploads.wait_for_actions(&actions).await;
1080 let withheld = manifest
1081 .predictions
1082 .iter()
1083 .filter(|prediction| {
1084 introduced.contains(&prediction.invocation)
1085 && unpublished.contains(&prediction.action)
1086 })
1087 .count();
1088 if withheld > 0 {
1089 warn!(
1092 "{withheld} of {} predicted actions were not published, so the remote task action manifest omits them",
1093 manifest.predictions.len()
1094 );
1095 manifest.predictions.retain(|prediction| {
1096 !(introduced.contains(&prediction.invocation)
1097 && unpublished.contains(&prediction.action))
1098 });
1099 }
1100 }
1101 match self
1102 .put_remote_task_manifest(&task, manifest, state.remote_etag)
1103 .await
1104 {
1105 Ok(remote_manifest) => {
1106 let _write_guard = self.manifest_write_lock.lock().unwrap();
1107 let reconciliation = (|| {
1108 let _file_guard = self.lock_task_manifest(&task)?;
1109 let manifest = match self.load_task_manifest(&task)? {
1110 Some(local) => {
1111 merge_remote_task_manifest(&task, remote_manifest, local).0
1112 }
1113 None => remote_manifest,
1114 };
1115 self.persist_task_manifest(&manifest)
1116 })();
1117 if let Err(error) = reconciliation {
1118 warn!(
1119 "remote task action manifest reconciliation failed for {task}: {error}"
1120 );
1121 }
1122 }
1123 Err(error) => {
1124 self.note_remote_failure();
1125 warn!("remote task action manifest upload failed for {task}: {error}");
1126 }
1127 }
1128 }
1129 Ok(completed)
1130 }
1131
1132 fn task_manifest_path(&self, task: &str) -> PathBuf {
1133 self.manifest_dir.join(format!("{task}.json"))
1134 }
1135
1136 fn task_manifest_lock_path(&self, task: &str) -> PathBuf {
1137 self.manifest_dir.join("locks").join(format!("{task}.lock"))
1138 }
1139
1140 fn lock_task_manifest(&self, task: &str) -> Result<fslock::LockFile> {
1141 let path = self.task_manifest_lock_path(task);
1142 fs::create_dir_all(path.parent().expect("task manifest lock has a parent"))?;
1143 let mut lock = fslock::LockFile::open(&path)?;
1144 lock.lock()?;
1145 Ok(lock)
1146 }
1147
1148 fn load_task_manifest(&self, task: &str) -> Result<Option<TaskActionManifest>> {
1149 match fs::read(self.task_manifest_path(task)) {
1150 Ok(contents) => Ok(Some(self.parse_task_manifest(task, &contents, false)?)),
1151 Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
1152 Err(error) => Err(error.into()),
1153 }
1154 }
1155
1156 fn parse_task_manifest(
1157 &self,
1158 task: &str,
1159 contents: &[u8],
1160 require_canonical: bool,
1161 ) -> Result<TaskActionManifest> {
1162 let manifest: TaskActionManifest = serde_json::from_slice(contents)?;
1163 validate_task_manifest(&manifest, task)?;
1164 if require_canonical && canonical_json(&manifest)? != contents {
1165 bail!("task action manifest is not canonical JSON");
1166 }
1167 Ok(manifest)
1168 }
1169
1170 fn task_manifest_selector(task: &str) -> Result<(Vec<u8>, CacheDigest)> {
1171 TaskActionManifest::selector(task)
1172 }
1173
1174 fn persist_task_manifest(&self, manifest: &TaskActionManifest) -> Result<()> {
1175 let bytes = canonical_json(manifest)?;
1176 fs::create_dir_all(self.manifest_dir.as_path())?;
1177 let mut temporary = tempfile::NamedTempFile::new_in(self.manifest_dir.as_path())?;
1178 std::io::Write::write_all(temporary.as_file_mut(), &bytes)?;
1179 temporary.as_file_mut().sync_all()?;
1180 temporary
1181 .persist(self.task_manifest_path(&manifest.task))
1182 .map_err(|error| error.error)?;
1183 Ok(())
1184 }
1185
1186 async fn get_remote_task_manifest(
1187 &self,
1188 task: &str,
1189 ) -> Result<Option<(TaskActionManifest, String)>> {
1190 let Some(remote) = &self.remote else {
1191 return Ok(None);
1192 };
1193 let (_, selector) = Self::task_manifest_selector(task)?;
1194 let _permit = self.remote_transfers.acquire().await?;
1195 self.stats
1196 .remote_manifest_lookups
1197 .fetch_add(1, Ordering::Relaxed);
1198 let _timer = AtomicDurationTimer::start(&self.stats.remote_manifest_lookup_duration_ns);
1199 let Some(remote_manifest) = remote.get_action_manifest(&selector).await? else {
1200 return Ok(None);
1201 };
1202 let manifest = self.parse_task_manifest(task, &remote_manifest.bytes, true)?;
1203 Ok(Some((manifest, remote_manifest.etag)))
1204 }
1205
1206 async fn put_remote_task_manifest(
1207 &self,
1208 task: &str,
1209 mut manifest: TaskActionManifest,
1210 mut expected_etag: Option<String>,
1211 ) -> Result<TaskActionManifest> {
1212 let Some(remote) = &self.remote else {
1213 return Ok(manifest);
1214 };
1215 let (_, selector) = Self::task_manifest_selector(task)?;
1216 for _ in 0..4 {
1217 let bytes = canonical_json(&manifest)?;
1218 let outcome = {
1219 let _permit = self.remote_transfers.acquire().await?;
1220 remote
1221 .put_action_manifest(&selector, &bytes, expected_etag.as_deref())
1222 .await?
1223 };
1224 match outcome {
1225 ManifestPutOutcome::Stored => return Ok(manifest),
1226 ManifestPutOutcome::PreconditionFailed => {
1227 let Some((remote_manifest, etag)) = self.get_remote_task_manifest(task).await?
1228 else {
1229 expected_etag = None;
1230 continue;
1231 };
1232 manifest = merge_task_manifests(task, Some(remote_manifest), manifest)?;
1233 expected_etag = Some(etag);
1234 }
1235 }
1236 }
1237 bail!("remote task action manifest changed too frequently")
1238 }
1239
1240 pub fn stats(&self) -> AgentStats {
1242 AgentStats {
1243 session_duration_ns: 0,
1244 lookups: self.stats.lookups.load(Ordering::Relaxed),
1245 unconsulted: self.stats.unconsulted.load(Ordering::Relaxed),
1246 hits: self.stats.hits.load(Ordering::Relaxed),
1247 stores: self.stats.stores.load(Ordering::Relaxed),
1248 stored_bytes: self.stats.stored_bytes.load(Ordering::Relaxed),
1249 verifications: self.stats.verifications.load(Ordering::Relaxed),
1250 divergences: self.stats.divergences.load(Ordering::Relaxed),
1251 downloaded_bytes: self.stats.downloaded_bytes.load(Ordering::Relaxed),
1252 uploaded_bytes: self.stats.uploaded_bytes.load(Ordering::Relaxed),
1253 background_uploads: self.stats.background_uploads.load(Ordering::Relaxed),
1254 background_upload_failures: self
1255 .stats
1256 .background_upload_failures
1257 .load(Ordering::Relaxed),
1258 remote_blob_pack_uploads: self.stats.remote_blob_pack_uploads.load(Ordering::Relaxed),
1259 remote_blob_pack_upload_blobs: self
1260 .stats
1261 .remote_blob_pack_upload_blobs
1262 .load(Ordering::Relaxed),
1263 upload_drain_duration_ns: self.stats.upload_drain_duration_ns.load(Ordering::Relaxed),
1264 prefetched_actions: self.stats.prefetched_actions.load(Ordering::Relaxed),
1265 predictions_loaded: self.stats.predictions_loaded.load(Ordering::Relaxed),
1266 bypasses: self.stats.bypasses.lock().unwrap().clone(),
1267 avoided_compiler_duration_ns: self
1268 .stats
1269 .avoided_compiler_duration_ns
1270 .load(Ordering::Relaxed),
1271 compiler: self.stats.compiler.lock().unwrap().clone(),
1272 slow_compilations: self.stats.slow_compilations.lock().unwrap().clone(),
1273 remote_failures: self.stats.remote_failures.load(Ordering::Relaxed),
1274 remote_manifest_lookups: self.stats.remote_manifest_lookups.load(Ordering::Relaxed),
1275 remote_manifest_lookup_duration_ns: self
1276 .stats
1277 .remote_manifest_lookup_duration_ns
1278 .load(Ordering::Relaxed),
1279 remote_action_lookups: self.stats.remote_action_lookups.load(Ordering::Relaxed),
1280 remote_action_lookup_duration_ns: self
1281 .stats
1282 .remote_action_lookup_duration_ns
1283 .load(Ordering::Relaxed),
1284 remote_blob_requests: self.stats.remote_blob_requests.load(Ordering::Relaxed),
1285 remote_blob_pack_requests: self.stats.remote_blob_pack_requests.load(Ordering::Relaxed),
1286 remote_blob_pack_blobs: self.stats.remote_blob_pack_blobs.load(Ordering::Relaxed),
1287 remote_blob_transfer_duration_ns: self
1288 .stats
1289 .remote_blob_transfer_duration_ns
1290 .load(Ordering::Relaxed),
1291 local_cas_write_duration_ns: self
1292 .stats
1293 .local_cas_write_duration_ns
1294 .load(Ordering::Relaxed),
1295 prefetch_runs: self.stats.prefetch_runs.load(Ordering::Relaxed),
1296 prefetch_duration_ns: self.stats.prefetch_duration_ns.load(Ordering::Relaxed),
1297 materialization_duration_ns: self
1298 .stats
1299 .materialization_duration_ns
1300 .load(Ordering::Relaxed),
1301 restored_output_files: self.stats.restored_output_files.load(Ordering::Relaxed),
1302 restored_output_bytes: self.stats.restored_output_bytes.load(Ordering::Relaxed),
1303 reflinked_output_files: self.stats.reflinked_output_files.load(Ordering::Relaxed),
1304 reflinked_output_bytes: self.stats.reflinked_output_bytes.load(Ordering::Relaxed),
1305 copied_output_files: self.stats.copied_output_files.load(Ordering::Relaxed),
1306 copied_output_bytes: self.stats.copied_output_bytes.load(Ordering::Relaxed),
1307 reused_output_files: self.stats.reused_output_files.load(Ordering::Relaxed),
1308 reused_output_bytes: self.stats.reused_output_bytes.load(Ordering::Relaxed),
1309 }
1310 }
1311
1312 pub async fn handle_requests(
1319 &self,
1320 requests: impl IntoIterator<Item = AgentRequest>,
1321 ) -> Vec<AgentResponse> {
1322 let mut connection = ConnectionUploads::default();
1323 let mut responses = Vec::new();
1324 for request in requests {
1325 responses.push(self.respond_on(request, &mut connection).await);
1326 }
1327 responses
1328 }
1329
1330 fn write_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1331 Self::digest_lock(&self.write_locks, digest)
1332 }
1333
1334 fn action_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1335 Self::digest_lock(&self.action_locks, digest)
1336 }
1337
1338 fn digest_lock(
1339 locks: &Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>,
1340 digest: &CacheDigest,
1341 ) -> Arc<tokio::sync::Mutex<()>> {
1342 let mut locks = locks.lock().unwrap();
1343 locks.retain(|_, lock| lock.strong_count() > 0);
1344 if let Some(lock) = locks.get(digest).and_then(Weak::upgrade) {
1345 return lock;
1346 }
1347 let lock = Arc::new(tokio::sync::Mutex::new(()));
1348 locks.insert(digest.clone(), Arc::downgrade(&lock));
1349 lock
1350 }
1351
1352 #[cfg(test)]
1357 async fn respond(&self, request: AgentRequest) -> AgentResponse {
1358 self.respond_on(request, &mut ConnectionUploads::default())
1359 .await
1360 }
1361
1362 async fn respond_on(
1363 &self,
1364 request: AgentRequest,
1365 connection: &mut ConnectionUploads,
1366 ) -> AgentResponse {
1367 match &request {
1368 AgentRequest::FindActionPrediction { task, .. }
1369 | AgentRequest::RecordActionPrediction { task, .. } => {
1370 self.ensure_task_loaded(task).await;
1371 }
1372 _ => {}
1373 }
1374 let result = match request {
1375 AgentRequest::BeginTask { task } => self
1376 .begin_task(&task)
1377 .await
1378 .map(|run| AgentResponse::TaskBegun { run }),
1379 AgentRequest::CommitTask { run } => self
1380 .commit_task(&run)
1381 .await
1382 .map(|()| AgentResponse::TaskCommitted),
1383 AgentRequest::FindBlob { digest } => self.find_blob(&digest).await,
1384 AgentRequest::FindBlobs { digests } => self.find_blobs(digests).await,
1385 AgentRequest::StoreBlob { digest, source } => {
1386 self.store_blob(&digest, &source, connection).await
1387 }
1388 AgentRequest::FindActionResult { action } => {
1389 self.stats.lookups.fetch_add(1, Ordering::Relaxed);
1390 self.find_action_result(&action).await
1391 }
1392 AgentRequest::RecordActionHit {
1393 action,
1394 restore,
1395 crate_name,
1396 } => {
1397 let diagnostic = connection
1398 .take_action_diagnostic("hit", crate_name.as_deref())
1399 .map(|diagnostic| AgentEvent::ActionDiagnostic {
1400 outcome: "hit".into(),
1401 crate_name: crate_name.clone(),
1402 diagnostic,
1403 });
1404 self.record_action_hit(&action, restore, crate_name, diagnostic)
1405 }
1406 AgentRequest::RecordBypass { kind } => {
1407 *self
1408 .stats
1409 .bypasses
1410 .lock()
1411 .unwrap()
1412 .entry(kind.clone())
1413 .or_insert(0) += 1;
1414 self.emit(|| AgentEvent::Bypass { kind });
1415 Ok(AgentResponse::BypassRecorded)
1416 }
1417 AgentRequest::RecordUnconsulted => {
1418 self.stats.unconsulted.fetch_add(1, Ordering::Relaxed);
1419 self.emit(|| AgentEvent::Unconsulted);
1420 Ok(AgentResponse::UnconsultedRecorded)
1421 }
1422 AgentRequest::RecordWarning { message } => match parse_action_diagnostic(&message) {
1423 Some(Ok((outcome, crate_name, diagnostic))) => {
1424 connection.record_action_diagnostic(outcome, crate_name, diagnostic);
1425 Ok(AgentResponse::WarningRecorded)
1426 }
1427 Some(Err(error)) => Err(error),
1428 None => self.record_warning(message),
1429 },
1430 AgentRequest::FindFileDigests { scope, files } => self.find_file_digests(scope, files),
1431 AgentRequest::JoinActionPromise {
1432 adapter,
1433 invocation,
1434 } => self.join_action_promise(&adapter, &invocation).await,
1435 AgentRequest::CompleteActionPromise { claim, prediction } => {
1436 self.complete_action_promise(&claim, &prediction).await
1437 }
1438 AgentRequest::RecordFileDigests { scope, entries } => {
1439 self.record_file_digests(scope, entries)
1440 }
1441 AgentRequest::RecordCompilerInvocation {
1442 outcome,
1443 crate_name,
1444 duration_ns,
1445 } => {
1446 let diagnostic = connection
1447 .take_action_diagnostic(&outcome, crate_name.as_deref())
1448 .map(|diagnostic| AgentEvent::ActionDiagnostic {
1449 outcome: outcome.clone(),
1450 crate_name: crate_name.clone(),
1451 diagnostic,
1452 });
1453 self.record_compiler_invocation(
1454 &outcome,
1455 crate_name.as_deref(),
1456 duration_ns,
1457 diagnostic,
1458 )
1459 }
1460 AgentRequest::RecordActionVerification { matched, restore } => {
1461 self.record_materialization(restore);
1462 self.stats.verifications.fetch_add(1, Ordering::Relaxed);
1463 if !matched {
1464 self.stats.divergences.fetch_add(1, Ordering::Relaxed);
1465 }
1466 self.emit(|| AgentEvent::Verification { matched, restore });
1467 Ok(AgentResponse::ActionVerificationRecorded)
1468 }
1469 AgentRequest::StoreActionResult { result } => {
1470 self.store_action_result(&result, connection).await
1471 }
1472 AgentRequest::FindActionPrediction { task, invocation } => {
1473 self.find_action_prediction(&task, &invocation)
1474 }
1475 AgentRequest::RecordActionPrediction { task, prediction } => {
1476 self.record_action_prediction(&task, prediction)
1477 }
1478 AgentRequest::FindExecutableIdentity {
1479 executable,
1480 environment,
1481 } => self.find_executable_identity(executable, environment),
1482 AgentRequest::StoreExecutableIdentity {
1483 executable,
1484 environment,
1485 stdout,
1486 } => self.store_executable_identity(executable, environment, stdout),
1487 AgentRequest::Hello { .. } => {
1488 Err(eyre::eyre!("hello is only valid as the first request"))
1489 }
1490 };
1491 result.unwrap_or_else(|error| AgentResponse::Error {
1492 message: error.to_string(),
1493 })
1494 }
1495
1496 async fn find_blob(&self, digest: &CacheDigest) -> Result<AgentResponse> {
1497 if let Some(path) = self.find_verified_blob(digest)? {
1498 return Ok(AgentResponse::Blob { path: Some(path) });
1499 }
1500 if !self.remote_mode.reads() {
1501 return Ok(AgentResponse::Blob { path: None });
1502 }
1503 let Some(remote) = &self.remote else {
1504 return Ok(AgentResponse::Blob { path: None });
1505 };
1506 match self.fetch_remote_blob(remote, digest).await {
1507 Ok(path) => Ok(AgentResponse::Blob { path: Some(path) }),
1508 Err(error) => {
1509 warn!(
1510 "remote cache blob lookup failed for {}: {error}",
1511 digest.hash
1512 );
1513 Ok(AgentResponse::Blob { path: None })
1514 }
1515 }
1516 }
1517
1518 async fn find_blobs(&self, digests: Vec<CacheDigest>) -> Result<AgentResponse> {
1519 let mut paths = BTreeMap::new();
1520 let mut missing = Vec::new();
1521 for digest in &digests {
1522 match self.find_verified_blob(digest)? {
1523 Some(path) => {
1524 paths.insert(digest.clone(), path);
1525 }
1526 None => {
1527 missing.push(digest.clone());
1528 }
1529 }
1530 }
1531
1532 if !missing.is_empty()
1533 && self.remote_mode.reads()
1534 && let Some(remote) = &self.remote
1535 {
1536 paths.extend(self.fetch_remote_blobs(remote, missing, None).await);
1537 }
1538
1539 Ok(AgentResponse::Blobs {
1540 paths: digests
1541 .into_iter()
1542 .map(|digest| paths.get(&digest).cloned())
1543 .collect(),
1544 })
1545 }
1546
1547 async fn store_blob(
1548 &self,
1549 digest: &CacheDigest,
1550 source: &Path,
1551 connection: &mut ConnectionUploads,
1552 ) -> Result<AgentResponse> {
1553 let path = {
1554 let lock = self.write_lock(digest);
1555 let _guard = lock.lock().await;
1556 if let Some(path) = self.find_verified_blob(digest)? {
1557 path
1558 } else {
1559 let path = self.cas.store_file(digest, source)?;
1560 self.remember_verified_blob(digest, &path);
1561 self.stats.stores.fetch_add(1, Ordering::Relaxed);
1562 self.stats
1563 .stored_bytes
1564 .fetch_add(digest.size, Ordering::Relaxed);
1565 path
1566 }
1567 };
1568 if let Some(uploads) = &self.uploads {
1571 uploads.queue_blob(digest, path.clone(), connection);
1572 }
1573 Ok(AgentResponse::Stored { path })
1574 }
1575
1576 fn find_verified_blob(&self, digest: &CacheDigest) -> Result<Option<PathBuf>> {
1577 let remembered = self.verified_blobs.lock().unwrap().get(digest).cloned();
1578 if let Some(remembered) = remembered {
1579 if remembered.is_unchanged() {
1590 return Ok(Some(remembered.path));
1591 }
1592 self.verified_blobs.lock().unwrap().remove(digest);
1593 }
1594 let path = self.cas.find(digest)?;
1595 if let Some(path) = &path {
1596 self.remember_verified_blob(digest, path);
1597 }
1598 Ok(path)
1599 }
1600
1601 fn remember_verified_blob(&self, digest: &CacheDigest, path: &Path) {
1602 let Some(verified) = VerifiedBlob::describe(path) else {
1603 return;
1607 };
1608 self.verified_blobs
1609 .lock()
1610 .unwrap()
1611 .insert(digest.clone(), verified);
1612 }
1613
1614 async fn find_action_result(&self, action: &CacheDigest) -> Result<AgentResponse> {
1615 if let Some(result) = self.actions.find(action)? {
1616 return Ok(AgentResponse::ActionResult {
1617 result: Some(result),
1618 });
1619 }
1620 if !self.remote_mode.reads() {
1621 return Ok(AgentResponse::ActionResult { result: None });
1622 }
1623 let Some(remote) = &self.remote else {
1624 return Ok(AgentResponse::ActionResult { result: None });
1625 };
1626 let lock = self.action_lock(action);
1627 let _guard = lock.lock().await;
1628 if let Some(result) = self.actions.find(action)? {
1629 return Ok(AgentResponse::ActionResult {
1630 result: Some(result),
1631 });
1632 }
1633 if let Some(result) = self
1634 .pending_remote_actions
1635 .lock()
1636 .unwrap()
1637 .get(action)
1638 .cloned()
1639 {
1640 return Ok(AgentResponse::ActionResult {
1641 result: Some(result),
1642 });
1643 }
1644 let _permit = self.remote_transfers.acquire().await?;
1645 match self.get_remote_action_result(remote, action).await {
1646 Ok(Some(result)) => {
1647 self.pending_remote_actions
1648 .lock()
1649 .unwrap()
1650 .insert(action.clone(), result.clone());
1651 Ok(AgentResponse::ActionResult {
1652 result: Some(result),
1653 })
1654 }
1655 Ok(None) => Ok(AgentResponse::ActionResult { result: None }),
1656 Err(error) => {
1657 self.note_remote_failure();
1658 warn!(
1659 "remote cache action lookup failed for {}: {error}",
1660 action.hash
1661 );
1662 Ok(AgentResponse::ActionResult { result: None })
1663 }
1664 }
1665 }
1666
1667 async fn store_action_result(
1668 &self,
1669 result: &RemoteActionResult,
1670 connection: &ConnectionUploads,
1671 ) -> Result<AgentResponse> {
1672 let path = self.actions.store(result)?;
1673 if let Some(uploads) = &self.uploads {
1674 uploads.queue_action_result(result, connection);
1675 }
1676 Ok(AgentResponse::ActionStored { path })
1677 }
1678
1679 async fn join_action_promise(
1680 &self,
1681 adapter: &str,
1682 invocation: &CacheDigest,
1683 ) -> Result<AgentResponse> {
1684 if !self.remote_mode.reads() || !self.remote_mode.writes() {
1685 return Ok(AgentResponse::ActionPromise {
1686 claim: None,
1687 prediction: None,
1688 });
1689 }
1690 let Some(remote) = &self.remote else {
1691 return Ok(AgentResponse::ActionPromise {
1692 claim: None,
1693 prediction: None,
1694 });
1695 };
1696 let deadline = Instant::now() + ACTION_PROMISE_WAIT;
1697 loop {
1698 let _permit = self.remote_transfers.acquire().await?;
1699 let state = remote.join_action_promise(invocation, adapter).await;
1700 drop(_permit);
1701 match state {
1702 Ok(Some(ActionPromiseState::Claimed { claim })) => {
1703 return Ok(AgentResponse::ActionPromise {
1704 claim: Some(claim),
1705 prediction: None,
1706 });
1707 }
1708 Ok(Some(ActionPromiseState::Complete { prediction })) => {
1709 return Ok(AgentResponse::ActionPromise {
1710 claim: None,
1711 prediction: Some(prediction),
1712 });
1713 }
1714 Ok(Some(ActionPromiseState::Pending { retry_after_ms }))
1715 if Instant::now() < deadline =>
1716 {
1717 tokio::time::sleep(Duration::from_millis(retry_after_ms.clamp(10, 5_000)))
1718 .await;
1719 }
1720 Ok(Some(ActionPromiseState::Pending { .. }) | None) => {
1721 return Ok(AgentResponse::ActionPromise {
1722 claim: None,
1723 prediction: None,
1724 });
1725 }
1726 Ok(Some(_)) => {
1727 return Ok(AgentResponse::ActionPromise {
1728 claim: None,
1729 prediction: None,
1730 });
1731 }
1732 Err(error) => {
1733 self.note_remote_failure();
1734 warn!(
1735 "remote cache action promise failed for {}: {error}",
1736 invocation.hash
1737 );
1738 return Ok(AgentResponse::ActionPromise {
1739 claim: None,
1740 prediction: None,
1741 });
1742 }
1743 }
1744 }
1745 }
1746
1747 async fn complete_action_promise(
1748 &self,
1749 claim: &str,
1750 prediction: &ActionPrediction,
1751 ) -> Result<AgentResponse> {
1752 prediction.validate()?;
1753 let Some(remote) = &self.remote else {
1754 return Ok(AgentResponse::ActionPromiseCompleted);
1755 };
1756 let Some(uploads) = &self.uploads else {
1757 return Ok(AgentResponse::ActionPromiseCompleted);
1758 };
1759 if uploads
1760 .wait_for_actions(std::slice::from_ref(&prediction.action))
1761 .await
1762 .contains(&prediction.action)
1763 {
1764 return Ok(AgentResponse::ActionPromiseCompleted);
1767 }
1768 let completion = ActionPromiseCompletion {
1769 claim: claim.to_string(),
1770 prediction: prediction.clone(),
1771 };
1772 let _permit = self.remote_transfers.acquire().await?;
1773 if let Err(error) = remote
1774 .complete_action_promise(&prediction.invocation, &completion)
1775 .await
1776 {
1777 self.note_remote_failure();
1778 warn!(
1779 "remote cache action promise completion failed for {}: {error}",
1780 prediction.invocation.hash
1781 );
1782 }
1783 Ok(AgentResponse::ActionPromiseCompleted)
1784 }
1785
1786 fn note_remote_failure(&self) {
1788 self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
1789 }
1790
1791 async fn get_remote_action_result(
1792 &self,
1793 remote: &RemoteCacheClient,
1794 action: &CacheDigest,
1795 ) -> Result<Option<RemoteActionResult>> {
1796 self.stats
1797 .remote_action_lookups
1798 .fetch_add(1, Ordering::Relaxed);
1799 let _timer = AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
1800 remote.get_action_result(action).await
1801 }
1802
1803 fn record_action_hit(
1804 &self,
1805 action: &CacheDigest,
1806 restore: RestoreStats,
1807 crate_name: Option<String>,
1808 diagnostic: Option<AgentEvent>,
1809 ) -> Result<AgentResponse> {
1810 validate_crate_name(crate_name.as_deref())?;
1811 if self.actions.find(action)?.is_none() {
1812 let pending = self.pending_remote_actions.lock().unwrap().remove(action);
1813 if let Some(result) = pending {
1814 self.actions.store(&result)?;
1815 } else {
1816 bail!("cannot record a hit for a missing action result");
1817 }
1818 }
1819 self.record_restore(restore);
1820 self.stats.hits.fetch_add(1, Ordering::Relaxed);
1821 self.emit_action(
1822 diagnostic,
1823 AgentEvent::ActionHit {
1824 crate_name,
1825 restore,
1826 },
1827 );
1828 Ok(AgentResponse::ActionHitRecorded)
1829 }
1830
1831 fn record_restore(&self, restore: RestoreStats) {
1832 self.record_materialization(restore);
1833 atomic_saturating_add(
1834 &self.stats.avoided_compiler_duration_ns,
1835 restore.avoided_compiler_duration_ns,
1836 );
1837 atomic_saturating_add(&self.stats.restored_output_files, restore.output_files);
1838 atomic_saturating_add(&self.stats.restored_output_bytes, restore.output_bytes);
1839 atomic_saturating_add(
1840 &self.stats.reflinked_output_files,
1841 restore.reflinked_output_files,
1842 );
1843 atomic_saturating_add(
1844 &self.stats.reflinked_output_bytes,
1845 restore.reflinked_output_bytes,
1846 );
1847 atomic_saturating_add(&self.stats.copied_output_files, restore.copied_output_files);
1848 atomic_saturating_add(&self.stats.copied_output_bytes, restore.copied_output_bytes);
1849 atomic_saturating_add(&self.stats.reused_output_files, restore.reused_output_files);
1850 atomic_saturating_add(&self.stats.reused_output_bytes, restore.reused_output_bytes);
1851 }
1852
1853 fn record_compiler_invocation(
1854 &self,
1855 outcome: &str,
1856 crate_name: Option<&str>,
1857 duration_ns: u64,
1858 diagnostic: Option<AgentEvent>,
1859 ) -> Result<AgentResponse> {
1860 if !matches!(
1861 outcome,
1862 "miss" | "unconsulted" | "bypass" | "verification" | "incremental"
1863 ) {
1864 bail!("invalid compiler invocation outcome");
1865 }
1866 validate_crate_name(crate_name)?;
1867 let mut compiler = self.stats.compiler.lock().unwrap();
1868 let stats = compiler.entry(outcome.to_string()).or_default();
1869 stats.invocations = stats.invocations.saturating_add(1);
1870 stats.duration_ns = stats.duration_ns.saturating_add(duration_ns);
1871 drop(compiler);
1872 if outcome != "verification"
1873 && let Some(crate_name) = crate_name.filter(|name| !name.is_empty())
1874 {
1875 let mut slow = self.stats.slow_compilations.lock().unwrap();
1876 let duration = slow.entry(crate_name.to_string()).or_default();
1877 *duration = duration.saturating_add(duration_ns);
1878 }
1879 self.emit_action(
1880 diagnostic,
1881 AgentEvent::CompilerInvocation {
1882 outcome: outcome.to_string(),
1883 crate_name: crate_name.map(str::to_string),
1884 duration_ns,
1885 },
1886 );
1887 Ok(AgentResponse::CompilerInvocationRecorded)
1888 }
1889
1890 fn record_materialization(&self, restore: RestoreStats) {
1891 atomic_saturating_add(&self.stats.materialization_duration_ns, restore.duration_ns);
1892 }
1893
1894 fn record_warning(&self, message: String) -> Result<AgentResponse> {
1903 if message.is_empty()
1904 || message.len() > MAX_WARNING_BYTES
1905 || message.contains(['\n', '\r', '\0'])
1906 {
1907 bail!("invalid shim warning");
1908 }
1909 let mut warnings = self.warnings.lock().unwrap();
1910 if !warnings.contains(&message) && warnings.len() < MAX_WARNINGS {
1911 eprintln!("mbx[warning]: {message}");
1912 self.emit(|| AgentEvent::Warning {
1913 message: message.clone(),
1914 });
1915 warnings.insert(message);
1916 }
1917 Ok(AgentResponse::WarningRecorded)
1918 }
1919
1920 fn find_file_digests(
1927 &self,
1928 scope: FileDigestScope,
1929 files: Vec<FileIdentity>,
1930 ) -> Result<AgentResponse> {
1931 if files.len() > MAX_FILE_DIGEST_BATCH {
1932 bail!("too many file-digest lookups in one request");
1933 }
1934 let ledger = self.file_digests.lock().unwrap();
1935 let digests = files
1936 .into_iter()
1937 .map(|file| {
1938 let recorded = ledger.get(&(scope, file.path.clone()))?;
1939 (recorded.file == file).then(|| recorded.digest.clone())
1940 })
1941 .collect();
1942 Ok(AgentResponse::FileDigests { digests })
1943 }
1944
1945 fn record_file_digests(
1951 &self,
1952 scope: FileDigestScope,
1953 entries: Vec<RecordedFileDigest>,
1954 ) -> Result<AgentResponse> {
1955 if entries.len() > MAX_FILE_DIGEST_BATCH {
1956 bail!("too many file-digest records in one request");
1957 }
1958 for entry in &entries {
1959 validate_file_digest_record(entry)?;
1960 }
1961 let mut ledger = self.file_digests.lock().unwrap();
1962 for entry in entries {
1963 if ledger.len() >= MAX_FILE_DIGEST_ENTRIES
1964 && !ledger.contains_key(&(scope, entry.file.path.clone()))
1965 {
1966 break;
1967 }
1968 ledger.insert((scope, entry.file.path.clone()), entry);
1969 }
1970 Ok(AgentResponse::FileDigestsRecorded)
1971 }
1972
1973 pub fn seed_file_digests(&self, entries: Vec<(FileDigestScope, RecordedFileDigest)>) -> usize {
1982 self.seed_file_digests_with(|| entries)
1983 }
1984
1985 pub fn seed_file_digests_with(
1991 &self,
1992 load: impl FnOnce() -> Vec<(FileDigestScope, RecordedFileDigest)>,
1993 ) -> usize {
1994 let mut ledger = self.file_digests.lock().unwrap();
1995 let entries = load();
1996 let mut seeded = 0;
1997 for (scope, entry) in entries {
1998 if ledger.len() >= MAX_FILE_DIGEST_ENTRIES {
1999 break;
2000 }
2001 if validate_file_digest_record(&entry).is_err() {
2002 continue;
2003 }
2004 if let std::collections::btree_map::Entry::Vacant(slot) =
2005 ledger.entry((scope, entry.file.path.clone()))
2006 {
2007 slot.insert(entry);
2008 seeded += 1;
2009 }
2010 }
2011 seeded
2012 }
2013
2014 pub fn file_digests(&self) -> Vec<(FileDigestScope, RecordedFileDigest)> {
2017 self.file_digests
2018 .lock()
2019 .unwrap()
2020 .iter()
2021 .map(|((scope, _), entry)| (*scope, entry.clone()))
2022 .collect()
2023 }
2024
2025 fn find_action_prediction(
2026 &self,
2027 task: &str,
2028 invocation: &CacheDigest,
2029 ) -> Result<AgentResponse> {
2030 validate_task_identity(task)?;
2031 invocation.validate()?;
2032 let (prediction, prefetch) = {
2033 let mut tasks = self.task_actions.lock().unwrap();
2034 let Some(state) = tasks.get_mut(task) else {
2035 return Ok(AgentResponse::ActionPrediction { prediction: None });
2036 };
2037 activate_prediction_adapter(state, invocation)
2038 };
2039 if let Some(prefetch) = prefetch {
2040 self.spawn_prefetch_predictions(prefetch);
2041 }
2042 Ok(AgentResponse::ActionPrediction { prediction })
2043 }
2044
2045 fn record_action_prediction(
2046 &self,
2047 task: &str,
2048 prediction: ActionPrediction,
2049 ) -> Result<AgentResponse> {
2050 validate_task_identity(task)?;
2051 prediction.validate()?;
2052 let mut tasks = self.task_actions.lock().unwrap();
2053 let state = tasks.entry(task.to_string()).or_default();
2054 if !state.predictions.contains_key(&prediction.invocation)
2055 && state.predictions.len() >= MAX_TASK_ACTION_PREDICTIONS
2056 {
2057 bail!("task action manifest contains too many predictions");
2058 }
2059 state
2060 .predictions
2061 .insert(prediction.invocation.clone(), prediction.clone());
2062 state
2063 .pending_predictions
2064 .insert(prediction.invocation.clone(), prediction);
2065 Ok(AgentResponse::ActionPredictionRecorded)
2066 }
2067
2068 fn executable_identity_key(
2069 &self,
2070 executable: PathBuf,
2071 environment: BTreeMap<String, Option<String>>,
2072 ) -> Result<ExecutableIdentityKey> {
2073 if !environment.keys().all(|name| {
2078 matches!(
2079 name.as_str(),
2080 "MBX_FUSE_LD"
2081 | "RUSTUP_HOME"
2082 | "RUSTUP_TOOLCHAIN"
2083 | "SDKROOT"
2084 | "MACOSX_DEPLOYMENT_TARGET"
2085 | "LIB"
2086 | "UCRTVersion"
2087 | "UniversalCRTSdkDir"
2088 | "VCToolsInstallDir"
2089 | "VCToolsVersion"
2090 | "WindowsSdkDir"
2091 | "WindowsSDKVersion"
2092 )
2093 }) {
2094 bail!("executable identity contains an unsupported environment variable");
2095 }
2096 Ok(ExecutableIdentityKey {
2097 executable,
2098 environment,
2099 })
2100 }
2101
2102 fn find_executable_identity(
2103 &self,
2104 executable: PathBuf,
2105 environment: BTreeMap<String, Option<String>>,
2106 ) -> Result<AgentResponse> {
2107 let key = self.executable_identity_key(executable, environment)?;
2108 let stdout = self
2109 .executable_identities
2110 .lock()
2111 .unwrap()
2112 .get(&key)
2113 .cloned();
2114 Ok(AgentResponse::ExecutableIdentity { stdout })
2115 }
2116
2117 fn store_executable_identity(
2118 &self,
2119 executable: PathBuf,
2120 environment: BTreeMap<String, Option<String>>,
2121 stdout: Vec<u8>,
2122 ) -> Result<AgentResponse> {
2123 if stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE {
2124 bail!("executable identity exceeds {MAX_EXECUTABLE_IDENTITY_SIZE} bytes");
2125 }
2126 let key = self.executable_identity_key(executable, environment)?;
2127 let mut identities = self.executable_identities.lock().unwrap();
2128 let is_new = !identities.contains_key(&key);
2129 let previous_size = identities.get(&key).map_or(0, Vec::len);
2130 if is_new && identities.len() >= MAX_EXECUTABLE_IDENTITIES {
2131 bail!("executable identity cache contains too many entries");
2132 }
2133 let retained_bytes = identities.values().map(Vec::len).sum::<usize>();
2134 if retained_bytes - previous_size + stdout.len() > MAX_EXECUTABLE_IDENTITY_BYTES {
2135 bail!("executable identity cache contains too many bytes");
2136 }
2137 identities.insert(key, stdout.clone());
2138 Ok(AgentResponse::ExecutableIdentity {
2139 stdout: Some(stdout),
2140 })
2141 }
2142
2143 pub async fn handle_connection<S>(&self, stream: S) -> Result<()>
2145 where
2146 S: AsyncRead + AsyncWrite + Unpin,
2147 {
2148 let (reader, mut writer) = tokio::io::split(stream);
2149 let mut reader = BufReader::new(reader);
2150 let hello = read_request(&mut reader)
2151 .await?
2152 .ok_or_else(|| eyre::eyre!("connection closed before the agent handshake"))?;
2153 let request: AgentRequest = serde_json::from_str(&hello)?;
2154 match request {
2155 AgentRequest::Hello {
2156 protocol,
2157 client_version,
2158 } if protocol == AGENT_PROTOCOL_VERSION && client_version == self.version.as_ref() => {}
2159 AgentRequest::Hello { protocol, .. } if protocol != AGENT_PROTOCOL_VERSION => {
2160 send_response(
2161 &mut writer,
2162 &AgentResponse::Error {
2163 message: format!(
2164 "unsupported agent protocol {protocol}; expected {AGENT_PROTOCOL_VERSION}"
2165 ),
2166 },
2167 )
2168 .await?;
2169 return Ok(());
2170 }
2171 AgentRequest::Hello { client_version, .. } => {
2172 send_response(
2173 &mut writer,
2174 &AgentResponse::Error {
2175 message: format!(
2176 "cache client {client_version} does not match agent {}",
2177 self.version
2178 ),
2179 },
2180 )
2181 .await?;
2182 return Ok(());
2183 }
2184 _ => bail!("the first agent request must be hello"),
2185 }
2186 send_response(
2187 &mut writer,
2188 &AgentResponse::Hello {
2189 protocol: AGENT_PROTOCOL_VERSION,
2190 agent_version: self.version.to_string(),
2191 },
2192 )
2193 .await?;
2194
2195 let mut connection = ConnectionUploads::default();
2199 while let Some(line) = read_request(&mut reader).await? {
2200 let response = match serde_json::from_str(&line) {
2201 Ok(request) => self.respond_on(request, &mut connection).await,
2202 Err(error) => AgentResponse::Error {
2203 message: format!("invalid agent request: {error}"),
2204 },
2205 };
2206 send_response(&mut writer, &response).await?;
2207 }
2208 Ok(())
2209 }
2210}
2211
2212async fn read_request<R>(reader: &mut R) -> Result<Option<String>>
2218where
2219 R: AsyncBufRead + Unpin,
2220{
2221 let mut line = Vec::new();
2222 loop {
2223 let available = reader.fill_buf().await?;
2224 if available.is_empty() {
2225 break;
2226 }
2227 let (consumed, complete) = match available.iter().position(|byte| *byte == b'\n') {
2228 Some(index) => (index, true),
2229 None => (available.len(), false),
2230 };
2231 if line.len() + consumed > MAX_REQUEST_BYTES {
2232 bail!("agent request exceeded {MAX_REQUEST_BYTES} bytes");
2233 }
2234 line.extend_from_slice(&available[..consumed]);
2235 reader.consume(consumed + usize::from(complete));
2237 if complete {
2238 return Ok(Some(String::from_utf8(line)?));
2239 }
2240 }
2241 if line.is_empty() {
2242 Ok(None)
2243 } else {
2244 Ok(Some(String::from_utf8(line)?))
2245 }
2246}
2247
2248async fn send_response(
2249 writer: &mut (impl AsyncWrite + Unpin),
2250 response: &AgentResponse,
2251) -> Result<()> {
2252 let mut encoded = serde_json::to_vec(response)?;
2253 encoded.push(b'\n');
2254 writer.write_all(&encoded).await?;
2255 writer.flush().await?;
2256 Ok(())
2257}
2258
2259#[cfg(test)]
2260#[path = "agent_tests.rs"]
2261mod tests;