Skip to main content

mbx_cache_core/
agent.rs

1use crate::{
2    ActionPrediction, BlobPackLimits, BlobSource, BlobUpload, CacheDigest, CacheDirectory,
3    LocalActionCache, LocalCas, MAX_STAGED_BLOB_PACK_BYTES, MAX_STAGED_BLOB_PACK_ITEMS,
4    ManifestPutOutcome, RemoteActionResult, RemoteCacheClient, RemoteCacheMode, RustcMetadata,
5    TaskActionManifest, blob_pack_chunk, canonical_json,
6};
7use eyre::{Context, Result, bail};
8use futures_util::{FutureExt, StreamExt, future::BoxFuture, stream};
9use log::warn;
10use serde::{Deserialize, Serialize};
11use std::collections::BTreeMap;
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};
17use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
18
19const MAX_EXECUTABLE_IDENTITIES: usize = 64;
20const MAX_EXECUTABLE_IDENTITY_SIZE: usize = 64 * 1024;
21const MAX_EXECUTABLE_IDENTITY_BYTES: usize = 256 * 1024;
22const TASK_ACTION_MANIFEST_VERSION: u8 = 1;
23const MAX_TASK_ACTION_PREDICTIONS: usize = 16 * 1024;
24const MAX_REMOTE_TRANSFERS: usize = 64;
25const MAX_PREFETCH_TRANSFERS: usize = 48;
26const MAX_PREFETCH_ACTION_BATCH: usize = 256;
27const PREFETCH_ACTION_BATCH_DELAY: Duration = Duration::from_millis(5);
28const MAX_PREFETCH_DIRECTORY_OBJECTS: usize = 100_000;
29const MAX_PREFETCH_OBJECTS_PER_WAVE: usize = 100_000;
30const DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES: u64 = 5 * 1024 * 1024 * 1024;
31
32/// Remote action-cache access owned by one task session.
33pub struct AgentRemoteCache {
34    /// Remote protocol client used by the agent.
35    pub client: RemoteCacheClient,
36    /// Permitted remote read/write operations.
37    pub mode: RemoteCacheMode,
38    /// Directory used for verified downloads before CAS ingestion.
39    pub staging_dir: PathBuf,
40}
41
42/// Wire protocol version used between an in-process cache agent and its shims.
43pub const AGENT_PROTOCOL_VERSION: u8 = 2;
44/// Largest single protocol request the agent will read.
45///
46/// Requests are small JSON objects; the largest legitimate ones carry an output
47/// tree or a batch of digests, which stay far below this.
48const MAX_REQUEST_BYTES: usize = 16 * 1024 * 1024;
49
50/// A request accepted by the task-scoped cache agent.
51#[derive(Debug, Clone, Serialize, Deserialize)]
52#[serde(tag = "type", rename_all = "snake_case")]
53pub enum AgentRequest {
54    /// Negotiate protocol and application versions.
55    Hello {
56        /// Agent protocol version understood by the caller.
57        protocol: u8,
58        /// Human-readable mbx client version.
59        client_version: String,
60    },
61    /// Resolve a blob to a session-verified local CAS path.
62    FindBlob {
63        /// Blob to resolve.
64        digest: CacheDigest,
65    },
66    /// Resolve blobs to session-verified local CAS paths.
67    FindBlobs {
68        /// Blobs to resolve, preserving request order in the response.
69        digests: Vec<CacheDigest>,
70    },
71    /// Import a file into the local content-addressed store.
72    StoreBlob {
73        /// Digest the source must match.
74        digest: CacheDigest,
75        /// File to verify and import.
76        source: PathBuf,
77    },
78    /// Look up an action-result record.
79    FindActionResult {
80        /// Action digest to resolve.
81        action: CacheDigest,
82    },
83    /// Account for a successfully restored cache hit.
84    RecordActionHit {
85        /// Action that supplied the outputs.
86        action: CacheDigest,
87        /// Restoration work performed by the adapter.
88        restore: RestoreStats,
89    },
90    /// A compilation the adapter declined to cache, grouped by reason.
91    RecordBypass {
92        /// Stable, low-cardinality bypass-reason name.
93        kind: String,
94    },
95    /// A compilation the adapter could not look up, having no key to look up
96    /// with. Distinct from a bypass: these are cached once compiled.
97    RecordUnconsulted,
98    /// Account for one real compiler invocation performed by the adapter.
99    RecordCompilerInvocation {
100        /// Stable outcome category such as `miss`, `unconsulted`, or `bypass`.
101        outcome: String,
102        /// Compiler crate name, when the invocation supplied one.
103        crate_name: Option<String>,
104        /// Wall time spent running the compiler.
105        duration_ns: u64,
106    },
107    /// Account for a cache hit that was rebuilt for correctness verification.
108    RecordActionVerification {
109        /// Whether rebuilt and cached outputs matched.
110        matched: bool,
111        /// Restoration work performed before rebuilding.
112        restore: RestoreStats,
113    },
114    /// Store an action-result record locally and enqueue remote publication.
115    StoreActionResult {
116        /// Action-result record to store.
117        result: RemoteActionResult,
118    },
119    /// Find an earlier input prediction for a task and invocation.
120    FindActionPrediction {
121        /// Stable task identity.
122        task: String,
123        /// Digest of the compiler invocation without discovered inputs.
124        invocation: CacheDigest,
125    },
126    /// Record an input prediction after a successful compilation.
127    RecordActionPrediction {
128        /// Stable task identity.
129        task: String,
130        /// Adapter-owned prediction record.
131        prediction: ActionPrediction,
132    },
133    /// Find cached identity output for an executable and environment.
134    FindExecutableIdentity {
135        /// Executable whose identity command would run.
136        executable: PathBuf,
137        /// Environment variables affecting identity output.
138        environment: BTreeMap<String, Option<String>>,
139    },
140    /// Cache identity output for an executable and environment.
141    StoreExecutableIdentity {
142        /// Executable whose identity command ran.
143        executable: PathBuf,
144        /// Environment variables affecting identity output.
145        environment: BTreeMap<String, Option<String>>,
146        /// Captured identity-command standard output.
147        stdout: Vec<u8>,
148    },
149}
150
151/// Local output restoration work performed by one action-cache adapter hit.
152#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
153#[serde(deny_unknown_fields)]
154pub struct RestoreStats {
155    /// Cumulative time spent materializing and validating output files.
156    pub duration_ns: u64,
157    /// Compiler wall time recorded when this action was originally produced.
158    /// Zero means no timing hint was available.
159    pub avoided_compiler_duration_ns: u64,
160    /// Number of compiler output files restored.
161    pub output_files: u64,
162    /// Declared size of compiler output files restored.
163    pub output_bytes: u64,
164    /// Number of restored output files that share data blocks with the CAS.
165    pub reflinked_output_files: u64,
166    /// Declared size of restored outputs that share data blocks with the CAS.
167    pub reflinked_output_bytes: u64,
168    /// Number of restored output files that required a byte-for-byte copy.
169    pub copied_output_files: u64,
170    /// Declared size of restored outputs that required a byte-for-byte copy.
171    pub copied_output_bytes: u64,
172}
173
174/// A response returned by the task-scoped cache agent.
175#[derive(Debug, Clone, Serialize, Deserialize)]
176#[serde(tag = "type", rename_all = "snake_case")]
177pub enum AgentResponse {
178    /// Successful protocol negotiation.
179    Hello {
180        /// Agent protocol version.
181        protocol: u8,
182        /// Human-readable agent version.
183        agent_version: String,
184    },
185    /// A local CAS path already verified against the requested digest.
186    Blob {
187        /// Verified local path, or `None` on a cache miss.
188        path: Option<PathBuf>,
189    },
190    /// Local CAS paths already verified against the requested digests.
191    Blobs {
192        /// Verified local paths or misses, in request order.
193        paths: Vec<Option<PathBuf>>,
194    },
195    /// A blob was stored locally.
196    Stored {
197        /// Path of the stored object in the local CAS.
198        path: PathBuf,
199    },
200    /// Result of an action lookup.
201    ActionResult {
202        /// Validated action result, or `None` on a cache miss.
203        result: Option<RemoteActionResult>,
204    },
205    /// Hit statistics were updated.
206    ActionHitRecorded,
207    /// Verification statistics were updated.
208    ActionVerificationRecorded,
209    /// Bypass statistics were updated.
210    BypassRecorded,
211    /// Unconsulted-compilation statistics were updated.
212    UnconsultedRecorded,
213    /// Compiler invocation accounting was recorded.
214    CompilerInvocationRecorded,
215    /// An action result was stored.
216    ActionStored {
217        /// Path of the stored local action-result record.
218        path: PathBuf,
219    },
220    /// Result of an input-prediction lookup.
221    ActionPrediction {
222        /// Matching prediction, or `None` when none is known.
223        prediction: Option<ActionPrediction>,
224    },
225    /// An input prediction was recorded.
226    ActionPredictionRecorded,
227    /// Result of an executable-identity lookup.
228    ExecutableIdentity {
229        /// Captured output, or `None` when no identity is cached.
230        stdout: Option<Vec<u8>>,
231    },
232    /// The request failed without terminating the agent connection.
233    Error {
234        /// Human-readable failure description.
235        message: String,
236    },
237}
238
239/// Aggregate cache activity for one task session.
240///
241/// The agent produces these; nothing outside this crate has cause to build one.
242/// Saying so keeps a new counter from being a breaking change, which is what
243/// this type exists to accumulate -- reach for [`AgentStats::default`] and
244/// assign the fields a test needs.
245#[derive(Debug, Clone, Default, PartialEq, Eq)]
246#[non_exhaustive]
247pub struct AgentStats {
248    /// End-to-end lifetime of the task-scoped cache session.
249    pub session_duration_ns: u64,
250    /// Number of action-result lookups.
251    pub lookups: u64,
252    /// Compilations no action-result lookup was possible for, because no usable
253    /// action key was available.
254    ///
255    /// Counted separately from a miss, which is a lookup that found nothing.
256    /// Both compile, but only a miss says a lookup happened.
257    pub unconsulted: u64,
258    /// Number of lookups that found a valid local action result.
259    pub hits: u64,
260    /// Number of newly stored content-addressed objects.
261    pub stores: u64,
262    /// Total size of newly stored objects.
263    pub stored_bytes: u64,
264    /// Number of cache hits compiled again for qualification.
265    pub verifications: u64,
266    /// Number of qualification builds that diverged from the cached result.
267    pub divergences: u64,
268    /// CAS payload bytes downloaded from the remote cache.
269    pub downloaded_bytes: u64,
270    /// CAS payload bytes uploaded to the remote cache.
271    pub uploaded_bytes: u64,
272    /// Complete actions staged before an adapter requested them.
273    pub prefetched_actions: u64,
274    /// Compilations that were not cacheable, counted by reason.
275    pub bypasses: BTreeMap<String, u64>,
276    /// Estimated compiler time avoided by restored action hits.
277    pub avoided_compiler_duration_ns: u64,
278    /// Real compiler work performed in this session, grouped by outcome.
279    pub compiler: BTreeMap<String, CompilerStats>,
280    /// Cumulative real compiler time by crate name.
281    pub slow_compilations: BTreeMap<String, u64>,
282    /// Remote cache operations that failed and were degraded to a local result.
283    ///
284    /// A remote cache that cannot be reached, or that answers in a way this
285    /// client refuses, costs hit rate rather than correctness, so every one of
286    /// these is recovered from rather than raised. Counting them is what keeps a
287    /// remote that is failing every request from reading as one that merely had
288    /// nothing to offer.
289    pub remote_failures: u64,
290    /// Number of task manifest requests made to the remote cache.
291    pub remote_manifest_lookups: u64,
292    /// Cumulative time spent requesting remote task manifests.
293    pub remote_manifest_lookup_duration_ns: u64,
294    /// Number of action-result requests made to the remote cache.
295    pub remote_action_lookups: u64,
296    /// Cumulative time spent requesting remote action results.
297    pub remote_action_lookup_duration_ns: u64,
298    /// Number of blob requests made to the remote cache.
299    pub remote_blob_requests: u64,
300    /// Number of packed blob requests made to the remote cache.
301    pub remote_blob_pack_requests: u64,
302    /// Number of verified blobs received through packed responses.
303    pub remote_blob_pack_blobs: u64,
304    /// Cumulative time spent downloading and verifying remote blobs.
305    pub remote_blob_transfer_duration_ns: u64,
306    /// Cumulative time spent ingesting downloaded blobs into the local CAS.
307    pub local_cas_write_duration_ns: u64,
308    /// Number of speculative prefetch runs started for task manifests.
309    pub prefetch_runs: u64,
310    /// Cumulative wall time of speculative task-manifest prefetch runs.
311    pub prefetch_duration_ns: u64,
312    /// Cumulative time spent staging or materializing and validating cached outputs.
313    pub materialization_duration_ns: u64,
314    /// Number of compiler output files restored from action hits.
315    pub restored_output_files: u64,
316    /// Declared size of compiler output files restored from action hits.
317    pub restored_output_bytes: u64,
318    /// Number of restored output files materialized with filesystem reflinks.
319    pub reflinked_output_files: u64,
320    /// Declared size of outputs materialized with filesystem reflinks.
321    pub reflinked_output_bytes: u64,
322    /// Number of restored output files materialized by copying their bytes.
323    pub copied_output_files: u64,
324    /// Declared size of outputs materialized by copying their bytes.
325    pub copied_output_bytes: u64,
326}
327
328/// Count and cumulative wall time for one compiler-invocation outcome.
329///
330/// Non-exhaustive for the same reason as [`AgentStats`], which holds these:
331/// leaving the outer bag open does not help if describing an outcome in more
332/// detail still breaks the type inside it.
333#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
334#[non_exhaustive]
335pub struct CompilerStats {
336    /// Number of compiler invocations observed.
337    pub invocations: u64,
338    /// Cumulative wall time spent in those invocations.
339    pub duration_ns: u64,
340}
341
342impl CompilerStats {
343    /// The counts observed for one outcome.
344    pub fn new(invocations: u64, duration_ns: u64) -> Self {
345        Self {
346            invocations,
347            duration_ns,
348        }
349    }
350}
351
352#[derive(Default)]
353struct AtomicAgentStats {
354    lookups: AtomicU64,
355    unconsulted: AtomicU64,
356    hits: AtomicU64,
357    stores: AtomicU64,
358    stored_bytes: AtomicU64,
359    verifications: AtomicU64,
360    divergences: AtomicU64,
361    downloaded_bytes: AtomicU64,
362    uploaded_bytes: AtomicU64,
363    prefetched_actions: AtomicU64,
364    remote_failures: AtomicU64,
365    remote_manifest_lookups: AtomicU64,
366    remote_manifest_lookup_duration_ns: AtomicU64,
367    remote_action_lookups: AtomicU64,
368    remote_action_lookup_duration_ns: AtomicU64,
369    remote_blob_requests: AtomicU64,
370    remote_blob_pack_requests: AtomicU64,
371    remote_blob_pack_blobs: AtomicU64,
372    remote_blob_transfer_duration_ns: AtomicU64,
373    local_cas_write_duration_ns: AtomicU64,
374    prefetch_runs: AtomicU64,
375    prefetch_duration_ns: AtomicU64,
376    materialization_duration_ns: AtomicU64,
377    bypasses: Mutex<BTreeMap<String, u64>>,
378    avoided_compiler_duration_ns: AtomicU64,
379    compiler: Mutex<BTreeMap<String, CompilerStats>>,
380    slow_compilations: Mutex<BTreeMap<String, u64>>,
381    restored_output_files: AtomicU64,
382    restored_output_bytes: AtomicU64,
383    reflinked_output_files: AtomicU64,
384    reflinked_output_bytes: AtomicU64,
385    copied_output_files: AtomicU64,
386    copied_output_bytes: AtomicU64,
387}
388
389struct AtomicDurationTimer<'a> {
390    started: Instant,
391    target: &'a AtomicU64,
392}
393
394impl<'a> AtomicDurationTimer<'a> {
395    fn start(target: &'a AtomicU64) -> Self {
396        Self {
397            started: Instant::now(),
398            target,
399        }
400    }
401}
402
403impl Drop for AtomicDurationTimer<'_> {
404    fn drop(&mut self) {
405        atomic_saturating_add(self.target, duration_ns(self.started));
406    }
407}
408
409fn duration_ns(started: Instant) -> u64 {
410    started.elapsed().as_nanos().try_into().unwrap_or(u64::MAX)
411}
412
413fn atomic_saturating_add(target: &AtomicU64, value: u64) {
414    let _ = target.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
415        Some(current.saturating_add(value))
416    });
417}
418
419fn queue_prefetch_digest(
420    verified: &BTreeMap<CacheDigest, PathBuf>,
421    pending: &mut BTreeMap<CacheDigest, ()>,
422    digest: CacheDigest,
423) {
424    if verified.contains_key(&digest) || pending.contains_key(&digest) {
425        return;
426    }
427    pending.insert(digest, ());
428}
429
430fn queue_prefetch_directory(
431    seen: &BTreeMap<CacheDigest, ()>,
432    pending: &mut BTreeMap<CacheDigest, ()>,
433    digest: CacheDigest,
434    limit: usize,
435) -> bool {
436    if seen.contains_key(&digest) || pending.contains_key(&digest) {
437        return true;
438    }
439    if seen.len().saturating_add(pending.len()) >= limit {
440        return false;
441    }
442    pending.insert(digest, ());
443    true
444}
445
446/// Shared state for an agent hosted by the process that owns a build session.
447///
448/// Transport listeners deliberately live in the embedder so the session
449/// lifecycle owns them. This type only contains ecosystem-independent CAS and
450/// protocol logic.
451#[derive(Clone)]
452pub struct CacheAgent {
453    cas: LocalCas,
454    actions: LocalActionCache,
455    verified_blobs: Arc<Mutex<BTreeMap<CacheDigest, PathBuf>>>,
456    version: Arc<str>,
457    write_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
458    action_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
459    stats: Arc<AtomicAgentStats>,
460    executable_identities: Arc<Mutex<BTreeMap<ExecutableIdentityKey, Vec<u8>>>>,
461    manifest_dir: Arc<PathBuf>,
462    task_actions: Arc<Mutex<BTreeMap<String, TaskActionState>>>,
463    next_task_run: Arc<AtomicU64>,
464    manifest_write_lock: Arc<Mutex<()>>,
465    remote: Option<Arc<RemoteCacheClient>>,
466    remote_mode: RemoteCacheMode,
467    remote_staging_dir: Arc<PathBuf>,
468    remote_download_limit: u64,
469    remote_download_bytes: Arc<AtomicU64>,
470    pending_remote_actions: Arc<Mutex<BTreeMap<CacheDigest, RemoteActionResult>>>,
471    remote_transfers: Arc<tokio::sync::Semaphore>,
472    prefetch_transfers: Arc<tokio::sync::Semaphore>,
473    prefetch_tasks: Arc<Mutex<Vec<tokio::task::JoinHandle<()>>>>,
474}
475
476#[derive(Debug, Clone, Default)]
477struct TaskActionState {
478    manifest: String,
479    baseline_loaded: bool,
480    predictions: BTreeMap<CacheDigest, ActionPrediction>,
481    pending_predictions: BTreeMap<CacheDigest, ActionPrediction>,
482    remote_etag: Option<String>,
483}
484
485struct PrefetchedAction {
486    adapter: String,
487    result: RemoteActionResult,
488}
489
490struct RemoteDownloadReservation {
491    counter: Arc<AtomicU64>,
492    reserved: u64,
493    committed: bool,
494}
495
496impl RemoteDownloadReservation {
497    fn bytes(&self) -> u64 {
498        self.reserved
499    }
500
501    fn commit(mut self, bytes: u64) {
502        debug_assert!(bytes <= self.reserved);
503        self.counter
504            .fetch_sub(self.reserved.saturating_sub(bytes), Ordering::AcqRel);
505        self.committed = true;
506    }
507}
508
509impl Drop for RemoteDownloadReservation {
510    fn drop(&mut self) {
511        if !self.committed {
512            self.counter.fetch_sub(self.reserved, Ordering::AcqRel);
513        }
514    }
515}
516
517#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
518struct ExecutableIdentityKey {
519    executable: PathBuf,
520    environment: BTreeMap<String, Option<String>>,
521}
522
523impl CacheAgent {
524    /// Create an agent backed by the cache rooted at `cache_dir`.
525    pub fn new(cache_dir: impl Into<PathBuf>, version: impl Into<Arc<str>>) -> Self {
526        Self::build(cache_dir.into(), version.into(), None, 0)
527    }
528
529    /// Create an agent with local-first access to a remote action cache.
530    pub fn new_remote(
531        cache_dir: impl Into<PathBuf>,
532        version: impl Into<Arc<str>>,
533        remote: AgentRemoteCache,
534    ) -> Self {
535        Self::build(
536            cache_dir.into(),
537            version.into(),
538            Some(remote),
539            DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES,
540        )
541    }
542
543    /// Create a remote agent with a cumulative download budget for the session.
544    pub fn new_remote_with_download_limit(
545        cache_dir: impl Into<PathBuf>,
546        version: impl Into<Arc<str>>,
547        remote: AgentRemoteCache,
548        max_remote_download_bytes: u64,
549    ) -> Self {
550        Self::build(
551            cache_dir.into(),
552            version.into(),
553            Some(remote),
554            max_remote_download_bytes,
555        )
556    }
557
558    fn build(
559        cache_dir: PathBuf,
560        version: Arc<str>,
561        remote: Option<AgentRemoteCache>,
562        remote_download_limit: u64,
563    ) -> Self {
564        let remote_mode = remote
565            .as_ref()
566            .map_or(RemoteCacheMode::ReadOnly, |remote| remote.mode);
567        let remote_staging_dir = remote.as_ref().map_or_else(
568            || cache_dir.join("remote"),
569            |remote| remote.staging_dir.clone(),
570        );
571        let remote = remote.map(|remote| Arc::new(remote.client));
572        Self {
573            cas: LocalCas::new(cache_dir.clone()),
574            actions: LocalActionCache::new(cache_dir.clone()),
575            verified_blobs: Arc::new(Mutex::new(BTreeMap::new())),
576            version,
577            write_locks: Arc::new(Mutex::new(BTreeMap::new())),
578            action_locks: Arc::new(Mutex::new(BTreeMap::new())),
579            stats: Arc::new(AtomicAgentStats::default()),
580            executable_identities: Arc::new(Mutex::new(BTreeMap::new())),
581            manifest_dir: Arc::new(task_manifest_dir(&cache_dir)),
582            task_actions: Arc::new(Mutex::new(BTreeMap::new())),
583            next_task_run: Arc::new(AtomicU64::new(0)),
584            manifest_write_lock: Arc::new(Mutex::new(())),
585            remote,
586            remote_mode,
587            remote_staging_dir: Arc::new(remote_staging_dir),
588            remote_download_limit,
589            remote_download_bytes: Arc::new(AtomicU64::new(0)),
590            pending_remote_actions: Arc::new(Mutex::new(BTreeMap::new())),
591            remote_transfers: Arc::new(tokio::sync::Semaphore::new(MAX_REMOTE_TRANSFERS)),
592            prefetch_transfers: Arc::new(tokio::sync::Semaphore::new(MAX_PREFETCH_TRANSFERS)),
593            prefetch_tasks: Arc::new(Mutex::new(Vec::new())),
594        }
595    }
596
597    fn reserve_remote_download(&self, bytes: u64) -> Result<RemoteDownloadReservation> {
598        let mut current = self.remote_download_bytes.load(Ordering::Acquire);
599        loop {
600            let next = current
601                .checked_add(bytes)
602                .ok_or_else(|| eyre::eyre!("remote cache download budget overflowed"))?;
603            if next > self.remote_download_limit {
604                bail!(
605                    "remote cache download budget exceeded: {} bytes requested with {} of {} bytes already used",
606                    bytes,
607                    current,
608                    self.remote_download_limit
609                );
610            }
611            match self.remote_download_bytes.compare_exchange_weak(
612                current,
613                next,
614                Ordering::AcqRel,
615                Ordering::Acquire,
616            ) {
617                Ok(_) => {
618                    return Ok(RemoteDownloadReservation {
619                        counter: self.remote_download_bytes.clone(),
620                        reserved: bytes,
621                        committed: false,
622                    });
623                }
624                Err(observed) => current = observed,
625            }
626        }
627    }
628
629    fn reserve_remote_download_up_to(&self, requested: u64) -> Result<RemoteDownloadReservation> {
630        let mut current = self.remote_download_bytes.load(Ordering::Acquire);
631        loop {
632            let reserved = requested.min(self.remote_download_limit.saturating_sub(current));
633            let next = current + reserved;
634            match self.remote_download_bytes.compare_exchange_weak(
635                current,
636                next,
637                Ordering::AcqRel,
638                Ordering::Acquire,
639            ) {
640                Ok(_) => {
641                    return Ok(RemoteDownloadReservation {
642                        counter: self.remote_download_bytes.clone(),
643                        reserved,
644                        committed: false,
645                    });
646                }
647                Err(observed) => current = observed,
648            }
649        }
650    }
651
652    /// Load the last committed action manifest for a task into this session.
653    pub async fn begin_task(&self, task: &str) -> Result<String> {
654        self.begin_task_with_remote_errors(task, false).await
655    }
656
657    /// Load a task and finish its prefetch, surfacing remote lookup failures.
658    pub async fn prefetch_task(&self, task: &str) -> Result<String> {
659        let run = self.begin_task_with_remote_errors(task, true).await?;
660        self.wait_for_prefetches().await;
661        Ok(run)
662    }
663
664    async fn begin_task_with_remote_errors(&self, task: &str, strict: bool) -> Result<String> {
665        validate_task_identity(task)?;
666        let (remote_manifest, mut remote_etag) = if self.remote_mode.reads() {
667            match self.get_remote_task_manifest(task).await {
668                Ok(Some((manifest, etag))) => (Some(manifest), Some(etag)),
669                Ok(None) => (None, None),
670                Err(error) => {
671                    if strict {
672                        return Err(error).wrap_err_with(|| {
673                            format!("remote task action manifest lookup failed for {task}")
674                        });
675                    }
676                    self.note_remote_failure();
677                    warn!("remote task action manifest lookup failed for {task}: {error}");
678                    (None, None)
679                }
680            }
681        } else {
682            (None, None)
683        };
684        let manifest = {
685            let _write_guard = self.manifest_write_lock.lock().unwrap();
686            let _file_guard = self.lock_task_manifest(task)?;
687            let local_manifest = self.load_task_manifest(task)?;
688            let manifest = match (remote_manifest, local_manifest) {
689                (Some(remote), Some(local)) => {
690                    let (manifest, merged) = merge_remote_task_manifest(task, remote, local);
691                    if !merged {
692                        remote_etag = None;
693                    }
694                    Some(manifest)
695                }
696                (Some(remote), None) => Some(remote),
697                (None, local) => local,
698            };
699            if let Some(manifest) = &manifest {
700                self.persist_task_manifest(manifest)?;
701            }
702            manifest
703        };
704        let state = if let Some(manifest) = manifest {
705            TaskActionState {
706                manifest: task.to_string(),
707                baseline_loaded: true,
708                predictions: manifest
709                    .predictions
710                    .into_iter()
711                    .map(|prediction| (prediction.invocation.clone(), prediction))
712                    .collect(),
713                pending_predictions: BTreeMap::new(),
714                remote_etag,
715            }
716        } else {
717            TaskActionState {
718                manifest: task.to_string(),
719                baseline_loaded: true,
720                remote_etag,
721                ..TaskActionState::default()
722            }
723        };
724        let sequence = self.next_task_run.fetch_add(1, Ordering::Relaxed);
725        let run =
726            CacheDigest::blake3(format!("{task}\0{}\0{sequence}", std::process::id()).as_bytes())
727                .hash;
728        let predictions = state.predictions.values().cloned().collect();
729        self.task_actions.lock().unwrap().insert(run.clone(), state);
730        self.spawn_prefetch_predictions(predictions);
731        Ok(run)
732    }
733
734    /// Cancel speculative downloads before the owning session exits.
735    pub async fn cancel_prefetches(&self) {
736        let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
737        for task in &tasks {
738            task.abort();
739        }
740        for task in tasks {
741            if let Err(error) = task.await
742                && !error.is_cancelled()
743            {
744                warn!("remote action prefetch task failed: {error}");
745            }
746        }
747    }
748
749    async fn wait_for_prefetches(&self) {
750        let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
751        for task in tasks {
752            if let Err(error) = task.await {
753                warn!("remote action prefetch task failed: {error}");
754            }
755        }
756    }
757
758    /// Atomically publish the completed actions collected by a task run.
759    pub async fn commit_task(&self, run: &str) -> Result<()> {
760        validate_task_identity(run)?;
761        let state = self
762            .task_actions
763            .lock()
764            .unwrap()
765            .get(run)
766            .cloned()
767            .ok_or_else(|| eyre::eyre!("task action manifest baseline was not loaded"))?;
768        if !state.baseline_loaded {
769            bail!("task action manifest baseline was not loaded");
770        }
771        let task = state.manifest;
772        validate_task_identity(&task)?;
773        let manifest = {
774            let _write_guard = self.manifest_write_lock.lock().unwrap();
775            let _file_guard = self.lock_task_manifest(&task)?;
776            let mut predictions = self
777                .load_task_manifest(&task)?
778                .map(|manifest| {
779                    manifest
780                        .predictions
781                        .into_iter()
782                        .map(|prediction| (prediction.invocation.clone(), prediction))
783                        .collect::<BTreeMap<_, _>>()
784                })
785                .unwrap_or_default();
786            // Only publish predictions recorded by this run. `predictions`
787            // also contains the baseline loaded by `begin_task`; extending
788            // with that snapshot would overwrite newer entries committed by
789            // another agent process after this run began.
790            predictions.extend(state.pending_predictions);
791            let manifest = TaskActionManifest {
792                version: TASK_ACTION_MANIFEST_VERSION,
793                task: task.clone(),
794                predictions: predictions.into_values().collect(),
795            };
796            validate_task_manifest(&manifest, &task)?;
797            self.persist_task_manifest(&manifest)?;
798            manifest
799        };
800        self.task_actions.lock().unwrap().remove(run);
801        if self.remote_mode.writes() {
802            match self
803                .put_remote_task_manifest(&task, manifest, state.remote_etag)
804                .await
805            {
806                Ok(remote_manifest) => {
807                    let _write_guard = self.manifest_write_lock.lock().unwrap();
808                    let reconciliation = (|| {
809                        let _file_guard = self.lock_task_manifest(&task)?;
810                        let manifest = match self.load_task_manifest(&task)? {
811                            Some(local) => {
812                                merge_remote_task_manifest(&task, remote_manifest, local).0
813                            }
814                            None => remote_manifest,
815                        };
816                        self.persist_task_manifest(&manifest)
817                    })();
818                    if let Err(error) = reconciliation {
819                        warn!(
820                            "remote task action manifest reconciliation failed for {task}: {error}"
821                        );
822                    }
823                }
824                Err(error) => {
825                    self.note_remote_failure();
826                    warn!("remote task action manifest upload failed for {task}: {error}");
827                }
828            }
829        }
830        Ok(())
831    }
832
833    fn task_manifest_path(&self, task: &str) -> PathBuf {
834        self.manifest_dir.join(format!("{task}.json"))
835    }
836
837    fn task_manifest_lock_path(&self, task: &str) -> PathBuf {
838        self.manifest_dir.join("locks").join(format!("{task}.lock"))
839    }
840
841    fn lock_task_manifest(&self, task: &str) -> Result<fslock::LockFile> {
842        let path = self.task_manifest_lock_path(task);
843        fs::create_dir_all(path.parent().expect("task manifest lock has a parent"))?;
844        let mut lock = fslock::LockFile::open(&path)?;
845        lock.lock()?;
846        Ok(lock)
847    }
848
849    fn load_task_manifest(&self, task: &str) -> Result<Option<TaskActionManifest>> {
850        match fs::read(self.task_manifest_path(task)) {
851            Ok(contents) => Ok(Some(self.parse_task_manifest(task, &contents, false)?)),
852            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
853            Err(error) => Err(error.into()),
854        }
855    }
856
857    fn parse_task_manifest(
858        &self,
859        task: &str,
860        contents: &[u8],
861        require_canonical: bool,
862    ) -> Result<TaskActionManifest> {
863        let manifest: TaskActionManifest = serde_json::from_slice(contents)?;
864        validate_task_manifest(&manifest, task)?;
865        if require_canonical && canonical_json(&manifest)? != contents {
866            bail!("task action manifest is not canonical JSON");
867        }
868        Ok(manifest)
869    }
870
871    fn task_manifest_selector(task: &str) -> Result<(Vec<u8>, CacheDigest)> {
872        TaskActionManifest::selector(task)
873    }
874
875    fn persist_task_manifest(&self, manifest: &TaskActionManifest) -> Result<()> {
876        let bytes = canonical_json(manifest)?;
877        fs::create_dir_all(self.manifest_dir.as_path())?;
878        let mut temporary = tempfile::NamedTempFile::new_in(self.manifest_dir.as_path())?;
879        std::io::Write::write_all(temporary.as_file_mut(), &bytes)?;
880        temporary.as_file_mut().sync_all()?;
881        temporary
882            .persist(self.task_manifest_path(&manifest.task))
883            .map_err(|error| error.error)?;
884        Ok(())
885    }
886
887    async fn get_remote_task_manifest(
888        &self,
889        task: &str,
890    ) -> Result<Option<(TaskActionManifest, String)>> {
891        let Some(remote) = &self.remote else {
892            return Ok(None);
893        };
894        let (_, selector) = Self::task_manifest_selector(task)?;
895        let _permit = self.remote_transfers.acquire().await?;
896        self.stats
897            .remote_manifest_lookups
898            .fetch_add(1, Ordering::Relaxed);
899        let _timer = AtomicDurationTimer::start(&self.stats.remote_manifest_lookup_duration_ns);
900        let Some(remote_manifest) = remote.get_action_manifest(&selector).await? else {
901            return Ok(None);
902        };
903        let manifest = self.parse_task_manifest(task, &remote_manifest.bytes, true)?;
904        Ok(Some((manifest, remote_manifest.etag)))
905    }
906
907    async fn put_remote_task_manifest(
908        &self,
909        task: &str,
910        mut manifest: TaskActionManifest,
911        mut expected_etag: Option<String>,
912    ) -> Result<TaskActionManifest> {
913        let Some(remote) = &self.remote else {
914            return Ok(manifest);
915        };
916        let (_, selector) = Self::task_manifest_selector(task)?;
917        for _ in 0..4 {
918            let bytes = canonical_json(&manifest)?;
919            let outcome = {
920                let _permit = self.remote_transfers.acquire().await?;
921                remote
922                    .put_action_manifest(&selector, &bytes, expected_etag.as_deref())
923                    .await?
924            };
925            match outcome {
926                ManifestPutOutcome::Stored => return Ok(manifest),
927                ManifestPutOutcome::PreconditionFailed => {
928                    let Some((remote_manifest, etag)) = self.get_remote_task_manifest(task).await?
929                    else {
930                        expected_etag = None;
931                        continue;
932                    };
933                    manifest = merge_task_manifests(task, Some(remote_manifest), manifest)?;
934                    expected_etag = Some(etag);
935                }
936            }
937        }
938        bail!("remote task action manifest changed too frequently")
939    }
940
941    /// Return a snapshot of this session's cache activity.
942    pub fn stats(&self) -> AgentStats {
943        AgentStats {
944            session_duration_ns: 0,
945            lookups: self.stats.lookups.load(Ordering::Relaxed),
946            unconsulted: self.stats.unconsulted.load(Ordering::Relaxed),
947            hits: self.stats.hits.load(Ordering::Relaxed),
948            stores: self.stats.stores.load(Ordering::Relaxed),
949            stored_bytes: self.stats.stored_bytes.load(Ordering::Relaxed),
950            verifications: self.stats.verifications.load(Ordering::Relaxed),
951            divergences: self.stats.divergences.load(Ordering::Relaxed),
952            downloaded_bytes: self.stats.downloaded_bytes.load(Ordering::Relaxed),
953            uploaded_bytes: self.stats.uploaded_bytes.load(Ordering::Relaxed),
954            prefetched_actions: self.stats.prefetched_actions.load(Ordering::Relaxed),
955            bypasses: self.stats.bypasses.lock().unwrap().clone(),
956            avoided_compiler_duration_ns: self
957                .stats
958                .avoided_compiler_duration_ns
959                .load(Ordering::Relaxed),
960            compiler: self.stats.compiler.lock().unwrap().clone(),
961            slow_compilations: self.stats.slow_compilations.lock().unwrap().clone(),
962            remote_failures: self.stats.remote_failures.load(Ordering::Relaxed),
963            remote_manifest_lookups: self.stats.remote_manifest_lookups.load(Ordering::Relaxed),
964            remote_manifest_lookup_duration_ns: self
965                .stats
966                .remote_manifest_lookup_duration_ns
967                .load(Ordering::Relaxed),
968            remote_action_lookups: self.stats.remote_action_lookups.load(Ordering::Relaxed),
969            remote_action_lookup_duration_ns: self
970                .stats
971                .remote_action_lookup_duration_ns
972                .load(Ordering::Relaxed),
973            remote_blob_requests: self.stats.remote_blob_requests.load(Ordering::Relaxed),
974            remote_blob_pack_requests: self.stats.remote_blob_pack_requests.load(Ordering::Relaxed),
975            remote_blob_pack_blobs: self.stats.remote_blob_pack_blobs.load(Ordering::Relaxed),
976            remote_blob_transfer_duration_ns: self
977                .stats
978                .remote_blob_transfer_duration_ns
979                .load(Ordering::Relaxed),
980            local_cas_write_duration_ns: self
981                .stats
982                .local_cas_write_duration_ns
983                .load(Ordering::Relaxed),
984            prefetch_runs: self.stats.prefetch_runs.load(Ordering::Relaxed),
985            prefetch_duration_ns: self.stats.prefetch_duration_ns.load(Ordering::Relaxed),
986            materialization_duration_ns: self
987                .stats
988                .materialization_duration_ns
989                .load(Ordering::Relaxed),
990            restored_output_files: self.stats.restored_output_files.load(Ordering::Relaxed),
991            restored_output_bytes: self.stats.restored_output_bytes.load(Ordering::Relaxed),
992            reflinked_output_files: self.stats.reflinked_output_files.load(Ordering::Relaxed),
993            reflinked_output_bytes: self.stats.reflinked_output_bytes.load(Ordering::Relaxed),
994            copied_output_files: self.stats.copied_output_files.load(Ordering::Relaxed),
995            copied_output_bytes: self.stats.copied_output_bytes.load(Ordering::Relaxed),
996        }
997    }
998
999    /// Handle requests without a transport connection.
1000    ///
1001    /// Persistent wrappers use this entry point when Cargo invokes mbx outside
1002    /// an orchestrated session. It intentionally has the same response
1003    /// semantics as [`Self::handle_connection`], while leaving framing and the
1004    /// version handshake to callers that actually cross a process boundary.
1005    pub async fn handle_requests(
1006        &self,
1007        requests: impl IntoIterator<Item = AgentRequest>,
1008    ) -> Vec<AgentResponse> {
1009        let mut responses = Vec::new();
1010        for request in requests {
1011            responses.push(self.respond(request).await);
1012        }
1013        responses
1014    }
1015
1016    fn write_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1017        Self::digest_lock(&self.write_locks, digest)
1018    }
1019
1020    fn action_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1021        Self::digest_lock(&self.action_locks, digest)
1022    }
1023
1024    fn digest_lock(
1025        locks: &Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>,
1026        digest: &CacheDigest,
1027    ) -> Arc<tokio::sync::Mutex<()>> {
1028        let mut locks = locks.lock().unwrap();
1029        locks.retain(|_, lock| lock.strong_count() > 0);
1030        if let Some(lock) = locks.get(digest).and_then(Weak::upgrade) {
1031            return lock;
1032        }
1033        let lock = Arc::new(tokio::sync::Mutex::new(()));
1034        locks.insert(digest.clone(), Arc::downgrade(&lock));
1035        lock
1036    }
1037
1038    fn spawn_prefetch_predictions(&self, predictions: Vec<ActionPrediction>) {
1039        if predictions.is_empty() || !self.remote_mode.reads() || self.remote.is_none() {
1040            return;
1041        }
1042        let agent = self.clone();
1043        let task = tokio::spawn(async move {
1044            agent.prefetch_predictions(predictions.iter()).await;
1045        });
1046        self.prefetch_tasks.lock().unwrap().push(task);
1047    }
1048
1049    async fn prefetch_predictions<'a>(
1050        &self,
1051        predictions: impl Iterator<Item = &'a ActionPrediction>,
1052    ) {
1053        if !self.remote_mode.reads() || self.remote.is_none() {
1054            return;
1055        }
1056        self.stats.prefetch_runs.fetch_add(1, Ordering::Relaxed);
1057        let _timer = AtomicDurationTimer::start(&self.stats.prefetch_duration_ns);
1058        let mut actions = BTreeMap::new();
1059        for prediction in predictions {
1060            actions
1061                .entry(prediction.action.clone())
1062                .or_insert_with(|| prediction.adapter.clone());
1063        }
1064        let mut actions = actions.into_iter();
1065        let mut tasks = tokio::task::JoinSet::new();
1066        for _ in 0..MAX_PREFETCH_TRANSFERS {
1067            let Some((action, adapter)) = actions.next() else {
1068                break;
1069            };
1070            let agent = self.clone();
1071            tasks.spawn(async move { agent.resolve_prefetch_action(action, adapter).await });
1072        }
1073        let mut resolved = Vec::new();
1074        while !tasks.is_empty() {
1075            let result = if resolved.is_empty() {
1076                tasks.join_next().await
1077            } else {
1078                match tokio::time::timeout(PREFETCH_ACTION_BATCH_DELAY, tasks.join_next()).await {
1079                    Ok(result) => result,
1080                    Err(_) => {
1081                        self.prefetch_resolved_actions(std::mem::take(&mut resolved))
1082                            .await;
1083                        continue;
1084                    }
1085                }
1086            };
1087            let Some(result) = result else {
1088                break;
1089            };
1090            match result {
1091                Ok(Ok(Some(action))) => resolved.push(action),
1092                Ok(Ok(None)) => {}
1093                Ok(Err(error)) => {
1094                    self.note_remote_failure();
1095                    warn!("remote action prefetch failed: {error}");
1096                }
1097                Err(error) => {
1098                    self.note_remote_failure();
1099                    warn!("remote action prefetch task failed: {error}");
1100                }
1101            }
1102            if let Some((action, adapter)) = actions.next() {
1103                let agent = self.clone();
1104                tasks.spawn(async move { agent.resolve_prefetch_action(action, adapter).await });
1105            }
1106            if resolved.len() == MAX_PREFETCH_ACTION_BATCH {
1107                self.prefetch_resolved_actions(std::mem::take(&mut resolved))
1108                    .await;
1109            }
1110        }
1111        if !resolved.is_empty() {
1112            self.prefetch_resolved_actions(resolved).await;
1113        }
1114    }
1115
1116    #[cfg(test)]
1117    async fn prefetch_action(&self, action: CacheDigest, adapter: String) -> Result<()> {
1118        if let Some(action) = self.resolve_prefetch_action(action, adapter).await? {
1119            self.prefetch_resolved_actions(vec![action]).await;
1120        }
1121        Ok(())
1122    }
1123
1124    async fn resolve_prefetch_action(
1125        &self,
1126        action: CacheDigest,
1127        adapter: String,
1128    ) -> Result<Option<PrefetchedAction>> {
1129        let remote = self
1130            .remote
1131            .as_ref()
1132            .ok_or_else(|| eyre::eyre!("remote cache is not configured"))?;
1133        let result = {
1134            let lock = self.action_lock(&action);
1135            let _guard = lock.lock().await;
1136            if self.actions.find(&action)?.is_some() {
1137                return Ok(None);
1138            }
1139            if let Some(result) = self
1140                .pending_remote_actions
1141                .lock()
1142                .unwrap()
1143                .get(&action)
1144                .cloned()
1145            {
1146                result
1147            } else {
1148                let _prefetch_permit = self.prefetch_transfers.acquire().await?;
1149                let result = {
1150                    let _permit = self.remote_transfers.acquire().await?;
1151                    self.get_remote_action_result(remote, &action).await?
1152                };
1153                let Some(result) = result else {
1154                    return Ok(None);
1155                };
1156                self.pending_remote_actions
1157                    .lock()
1158                    .unwrap()
1159                    .insert(action.clone(), result.clone());
1160                result
1161            }
1162        };
1163        Ok(Some(PrefetchedAction { adapter, result }))
1164    }
1165
1166    fn prefetch_resolved_actions(&self, actions: Vec<PrefetchedAction>) -> BoxFuture<'_, ()> {
1167        self.prefetch_resolved_actions_inner(actions).boxed()
1168    }
1169
1170    async fn prefetch_resolved_actions_inner(&self, actions: Vec<PrefetchedAction>) {
1171        let Some(remote) = self.remote.as_deref() else {
1172            return;
1173        };
1174        if actions.is_empty() {
1175            return;
1176        }
1177
1178        let mut top_level = BTreeMap::new();
1179        for action in &actions {
1180            for digest in [
1181                Some(&action.result.action),
1182                action.result.metadata.as_ref(),
1183                action.result.output_root.as_ref(),
1184            ]
1185            .into_iter()
1186            .flatten()
1187            {
1188                top_level.insert(digest.clone(), ());
1189            }
1190        }
1191        let mut verified = self
1192            .fetch_remote_blobs(
1193                remote,
1194                top_level.into_keys().collect(),
1195                Some(&self.prefetch_transfers),
1196            )
1197            .await;
1198
1199        let mut next = BTreeMap::new();
1200        let mut pending_directories = BTreeMap::new();
1201        let mut parsed_directories = BTreeMap::new();
1202        let mut rustc_metadata = BTreeMap::new();
1203        for action in &actions {
1204            if action.adapter == "rustc"
1205                && let Some(metadata_digest) = &action.result.metadata
1206            {
1207                match verified
1208                    .get(metadata_digest)
1209                    .ok_or_else(|| eyre::eyre!("remote rustc action metadata is missing"))
1210                    .and_then(|path| Self::parse_rustc_metadata(path))
1211                {
1212                    Ok(metadata) => {
1213                        queue_prefetch_digest(&verified, &mut next, metadata.stdout.clone());
1214                        queue_prefetch_digest(&verified, &mut next, metadata.stderr.clone());
1215                        rustc_metadata.insert(metadata_digest.clone(), metadata);
1216                    }
1217                    Err(error) => warn!(
1218                        "remote rustc action metadata prefetch failed for {}: {error}",
1219                        action.result.action.hash
1220                    ),
1221                }
1222            }
1223            if let Some(output_root) = &action.result.output_root {
1224                pending_directories.insert(output_root.clone(), ());
1225            }
1226        }
1227
1228        let mut seen_directories = BTreeMap::new();
1229        loop {
1230            let mut following = BTreeMap::new();
1231            let mut directory_limit_exceeded = false;
1232            for digest in pending_directories.into_keys() {
1233                following.remove(&digest);
1234                if seen_directories.insert(digest.clone(), ()).is_some() {
1235                    continue;
1236                }
1237                if seen_directories.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
1238                    warn!("remote action output tree is too large to prefetch");
1239                    following.clear();
1240                    break;
1241                }
1242                match verified
1243                    .get(&digest)
1244                    .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))
1245                    .and_then(|path| Self::parse_cache_directory(path))
1246                {
1247                    Ok(directory) => {
1248                        for file in &directory.files {
1249                            queue_prefetch_digest(&verified, &mut next, file.digest.clone());
1250                            if next.len() >= MAX_PREFETCH_OBJECTS_PER_WAVE {
1251                                self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1252                                    .await;
1253                            }
1254                        }
1255                        for child in &directory.directories {
1256                            if !queue_prefetch_directory(
1257                                &seen_directories,
1258                                &mut following,
1259                                child.digest.clone(),
1260                                MAX_PREFETCH_DIRECTORY_OBJECTS,
1261                            ) {
1262                                warn!("remote action output tree is too large to prefetch");
1263                                directory_limit_exceeded = true;
1264                                break;
1265                            }
1266                            queue_prefetch_digest(&verified, &mut next, child.digest.clone());
1267                            if next.len() >= MAX_PREFETCH_OBJECTS_PER_WAVE {
1268                                self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1269                                    .await;
1270                            }
1271                        }
1272                        parsed_directories.insert(digest, directory);
1273                    }
1274                    Err(error) => warn!(
1275                        "remote action output directory prefetch failed for {}: {error}",
1276                        digest.hash
1277                    ),
1278                }
1279                if directory_limit_exceeded {
1280                    following.clear();
1281                    break;
1282                }
1283            }
1284            self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1285                .await;
1286            if following.is_empty() {
1287                break;
1288            }
1289            pending_directories = following;
1290        }
1291
1292        for action in actions {
1293            match Self::validate_prefetched_action(
1294                &action,
1295                &verified,
1296                &rustc_metadata,
1297                &parsed_directories,
1298            ) {
1299                Ok(()) => {
1300                    if let Err(error) = self.actions.store(&action.result) {
1301                        warn!(
1302                            "remote action prefetch could not publish {}: {error}",
1303                            action.result.action.hash
1304                        );
1305                        continue;
1306                    }
1307                    self.pending_remote_actions
1308                        .lock()
1309                        .unwrap()
1310                        .remove(&action.result.action);
1311                    self.stats
1312                        .prefetched_actions
1313                        .fetch_add(1, Ordering::Relaxed);
1314                }
1315                Err(error) => warn!(
1316                    "remote action prefetch was incomplete for {}: {error}",
1317                    action.result.action.hash
1318                ),
1319            }
1320        }
1321    }
1322
1323    async fn flush_prefetch_digest_batch(
1324        &self,
1325        remote: &RemoteCacheClient,
1326        verified: &mut BTreeMap<CacheDigest, PathBuf>,
1327        pending: &mut BTreeMap<CacheDigest, ()>,
1328    ) {
1329        if pending.is_empty() {
1330            return;
1331        }
1332        let digests = std::mem::take(pending).into_keys().collect();
1333        verified.extend(
1334            self.fetch_remote_blobs(remote, digests, Some(&self.prefetch_transfers))
1335                .await,
1336        );
1337    }
1338
1339    async fn fetch_remote_blobs(
1340        &self,
1341        remote: &RemoteCacheClient,
1342        digests: Vec<CacheDigest>,
1343        prefetch_limit: Option<&tokio::sync::Semaphore>,
1344    ) -> BTreeMap<CacheDigest, PathBuf> {
1345        let mut verified = BTreeMap::new();
1346        let mut missing = BTreeMap::new();
1347        for digest in digests {
1348            match self.find_verified_blob(&digest) {
1349                Ok(Some(path)) => {
1350                    verified.insert(digest, path);
1351                }
1352                Ok(None) => {
1353                    missing.insert(digest, ());
1354                }
1355                Err(error) => warn!(
1356                    "local cache blob lookup failed for {}: {error}",
1357                    digest.hash
1358                ),
1359            }
1360        }
1361        if missing.is_empty() {
1362            return verified;
1363        }
1364
1365        let mut pack_candidates = missing.clone();
1366        while !pack_candidates.is_empty() {
1367            let candidates = match blob_pack_chunk(
1368                &pack_candidates.keys().cloned().collect::<Vec<_>>(),
1369                BlobPackLimits {
1370                    max_items: MAX_STAGED_BLOB_PACK_ITEMS,
1371                    max_bytes: MAX_STAGED_BLOB_PACK_BYTES,
1372                },
1373            ) {
1374                Ok(candidates) if !candidates.is_empty() => candidates,
1375                Ok(_) => break,
1376                Err(error) => {
1377                    warn!("remote cache blob pack skipped: {error}");
1378                    break;
1379                }
1380            };
1381            // A pack and an individual fetch share these per-digest locks. Hold
1382            // them through ingestion so overlapping prefetch and foreground
1383            // requests cannot download or charge the same object twice.
1384            let mut pack_guards = BTreeMap::new();
1385            for digest in candidates {
1386                let guard = self.write_lock(&digest).lock_owned().await;
1387                match self.find_verified_blob(&digest) {
1388                    Ok(Some(path)) => {
1389                        pack_candidates.remove(&digest);
1390                        missing.remove(&digest);
1391                        verified.insert(digest, path);
1392                    }
1393                    Ok(None) => {
1394                        pack_guards.insert(digest, guard);
1395                    }
1396                    Err(error) => {
1397                        warn!(
1398                            "local cache blob lookup failed for {}: {error}",
1399                            digest.hash
1400                        );
1401                        pack_guards.insert(digest, guard);
1402                    }
1403                }
1404            }
1405            let requested = pack_guards.keys().cloned().collect::<Vec<_>>();
1406            if requested.is_empty() {
1407                continue;
1408            }
1409            let requested_bytes = requested
1410                .iter()
1411                .fold(0_u64, |total, digest| total.saturating_add(digest.size));
1412            let pack_reservation = match self
1413                .reserve_remote_download_up_to(requested_bytes.min(MAX_STAGED_BLOB_PACK_BYTES))
1414            {
1415                Ok(reservation) if reservation.bytes() > 0 => reservation,
1416                Ok(_) => break,
1417                Err(error) => {
1418                    warn!("remote cache blob pack skipped: {error}");
1419                    break;
1420                }
1421            };
1422            let (pack, transfer_duration_ns) = {
1423                let _prefetch_permit = match prefetch_limit {
1424                    Some(limit) => match limit.acquire().await {
1425                        Ok(permit) => Some(permit),
1426                        Err(error) => {
1427                            warn!(
1428                                "remote cache blob pack could not acquire prefetch limit: {error}"
1429                            );
1430                            break;
1431                        }
1432                    },
1433                    None => None,
1434                };
1435                let _transfer_permit = match self.remote_transfers.acquire().await {
1436                    Ok(permit) => permit,
1437                    Err(error) => {
1438                        warn!("remote cache blob pack could not acquire transfer limit: {error}");
1439                        break;
1440                    }
1441                };
1442                let transfer_started = Instant::now();
1443                let pack = remote
1444                    .get_blob_pack_with_limit(
1445                        &requested,
1446                        self.remote_staging_dir.as_path(),
1447                        pack_reservation.bytes(),
1448                    )
1449                    .await;
1450                (pack, duration_ns(transfer_started))
1451            };
1452            let pack = match pack {
1453                Ok(Some(pack)) => pack,
1454                Ok(None) => break,
1455                Err(error) => {
1456                    atomic_saturating_add(
1457                        &self.stats.remote_blob_transfer_duration_ns,
1458                        transfer_duration_ns,
1459                    );
1460                    warn!(
1461                        "remote cache blob pack failed; falling back to individual blobs: {error}"
1462                    );
1463                    break;
1464                }
1465            };
1466            pack_reservation.commit(pack.payload_bytes);
1467            atomic_saturating_add(
1468                &self.stats.remote_blob_transfer_duration_ns,
1469                transfer_duration_ns,
1470            );
1471            atomic_saturating_add(&self.stats.remote_blob_pack_requests, pack.requests);
1472            atomic_saturating_add(&self.stats.remote_blob_pack_blobs, pack.blob_count);
1473            atomic_saturating_add(&self.stats.downloaded_bytes, pack.payload_bytes);
1474            if pack.requested.is_empty() {
1475                // The server's negotiated cap can be smaller than the local
1476                // staging cap used to select this locked slice. Fall back for
1477                // this slice, but keep packing later candidates that may fit.
1478                for digest in &requested {
1479                    pack_candidates.remove(digest);
1480                }
1481                continue;
1482            }
1483            for digest in &pack.requested {
1484                pack_candidates.remove(digest);
1485            }
1486            let mut ingests = stream::iter(pack.blobs.into_iter().map(|(digest, source)| {
1487                let digest_for_result = digest.clone();
1488                let guard = pack_guards
1489                    .remove(&digest)
1490                    .expect("requested packed blob has a write lock");
1491                async move {
1492                    (
1493                        digest_for_result,
1494                        self.ingest_packed_blob(digest, source, guard).await,
1495                    )
1496                }
1497            }))
1498            .buffer_unordered(MAX_PREFETCH_TRANSFERS);
1499            while let Some((digest, result)) = ingests.next().await {
1500                match result {
1501                    Ok(path) => {
1502                        missing.remove(&digest);
1503                        verified.insert(digest, path);
1504                    }
1505                    Err(error) => warn!(
1506                        "remote cache packed blob ingest failed for {}: {error}",
1507                        digest.hash
1508                    ),
1509                }
1510            }
1511        }
1512
1513        let mut transfers = stream::iter(missing.into_keys().map(|digest| {
1514            let digest_for_result = digest.clone();
1515            async move {
1516                (
1517                    digest_for_result,
1518                    self.fetch_remote_blob_with_limit(remote, &digest, prefetch_limit)
1519                        .await,
1520                )
1521            }
1522        }))
1523        .buffer_unordered(MAX_PREFETCH_TRANSFERS);
1524        while let Some((digest, result)) = transfers.next().await {
1525            match result {
1526                Ok(path) => {
1527                    verified.insert(digest, path);
1528                }
1529                Err(error) => warn!(
1530                    "remote cache blob prefetch failed for {}: {error}",
1531                    digest.hash
1532                ),
1533            }
1534        }
1535        verified
1536    }
1537
1538    async fn ingest_packed_blob(
1539        &self,
1540        digest: CacheDigest,
1541        source: PathBuf,
1542        _guard: tokio::sync::OwnedMutexGuard<()>,
1543    ) -> Result<PathBuf> {
1544        let digest_size = digest.size;
1545        let agent = self.clone();
1546        let (path, stored, cas_duration_ns) = tokio::task::spawn_blocking(move || {
1547            if let Some(path) = agent.find_verified_blob(&digest)? {
1548                return Ok::<_, eyre::Report>((path, false, 0));
1549            }
1550            let cas_started = Instant::now();
1551            let path = agent.cas.store_verified_file(&digest, &source)?;
1552            let cas_duration_ns = duration_ns(cas_started);
1553            agent.remember_verified_blob(&digest, &path);
1554            Ok((path, true, cas_duration_ns))
1555        })
1556        .await??;
1557        atomic_saturating_add(&self.stats.local_cas_write_duration_ns, cas_duration_ns);
1558        if stored {
1559            self.stats.stores.fetch_add(1, Ordering::Relaxed);
1560            atomic_saturating_add(&self.stats.stored_bytes, digest_size);
1561        }
1562        Ok(path)
1563    }
1564
1565    fn parse_rustc_metadata(path: &Path) -> Result<RustcMetadata> {
1566        let bytes = fs::read(path)?;
1567        let metadata: RustcMetadata = serde_json::from_slice(&bytes)?;
1568        if metadata.version != 1 || metadata.kind != "rustc" || canonical_json(&metadata)? != bytes
1569        {
1570            bail!("remote rustc action metadata is invalid");
1571        }
1572        Ok(metadata)
1573    }
1574
1575    fn parse_cache_directory(path: &Path) -> Result<CacheDirectory> {
1576        let bytes = fs::read(path)?;
1577        let directory: CacheDirectory = serde_json::from_slice(&bytes)?;
1578        if directory.version != 1 || canonical_json(&directory)? != bytes {
1579            bail!("remote action output directory is invalid");
1580        }
1581        Ok(directory)
1582    }
1583
1584    #[cfg(test)]
1585    fn load_cache_directory(&self, digest: &CacheDigest) -> Result<CacheDirectory> {
1586        let path = self
1587            .find_verified_blob(digest)?
1588            .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))?;
1589        Self::parse_cache_directory(&path)
1590    }
1591
1592    fn validate_prefetched_action(
1593        action: &PrefetchedAction,
1594        verified: &BTreeMap<CacheDigest, PathBuf>,
1595        rustc_metadata: &BTreeMap<CacheDigest, RustcMetadata>,
1596        directories: &BTreeMap<CacheDigest, CacheDirectory>,
1597    ) -> Result<()> {
1598        if !verified.contains_key(&action.result.action) {
1599            bail!("remote action descriptor is missing");
1600        }
1601        if let Some(metadata) = &action.result.metadata {
1602            if action.adapter == "rustc" {
1603                let metadata = rustc_metadata
1604                    .get(metadata)
1605                    .ok_or_else(|| eyre::eyre!("remote rustc action metadata is missing"))?;
1606                for digest in [&metadata.stdout, &metadata.stderr] {
1607                    if !verified.contains_key(digest) {
1608                        bail!("remote rustc action diagnostic blob is missing");
1609                    }
1610                }
1611            } else if !verified.contains_key(metadata) {
1612                bail!("remote action metadata is missing");
1613            }
1614        }
1615        let mut pending = action
1616            .result
1617            .output_root
1618            .iter()
1619            .cloned()
1620            .collect::<Vec<_>>();
1621        let mut seen = BTreeMap::new();
1622        while let Some(digest) = pending.pop() {
1623            if seen.insert(digest.clone(), ()).is_some() {
1624                continue;
1625            }
1626            if seen.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
1627                bail!("remote action output tree is too large");
1628            }
1629            let directory = directories
1630                .get(&digest)
1631                .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))?;
1632            for file in &directory.files {
1633                if !verified.contains_key(&file.digest) {
1634                    bail!("remote action output file is missing");
1635                }
1636            }
1637            pending.extend(
1638                directory
1639                    .directories
1640                    .iter()
1641                    .map(|directory| directory.digest.clone()),
1642            );
1643        }
1644        Ok(())
1645    }
1646
1647    #[cfg(test)]
1648    async fn prefetch_output_tree(
1649        &self,
1650        remote: &RemoteCacheClient,
1651        output_root: &CacheDigest,
1652    ) -> Result<()> {
1653        let mut pending = vec![output_root.clone()];
1654        let mut seen = BTreeMap::new();
1655        while let Some(digest) = pending.pop() {
1656            if seen.insert(digest.clone(), ()).is_some() {
1657                continue;
1658            }
1659            if seen.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
1660                bail!("remote action output tree is too large");
1661            }
1662            self.fetch_remote_blob_with_limit(remote, &digest, Some(&self.prefetch_transfers))
1663                .await?;
1664            let directory = self.load_cache_directory(&digest)?;
1665            let mut transfers = stream::iter(directory.files.into_iter().map(|file| async move {
1666                self.fetch_remote_blob_with_limit(
1667                    remote,
1668                    &file.digest,
1669                    Some(&self.prefetch_transfers),
1670                )
1671                .await
1672                .map(|_| ())
1673            }))
1674            .buffer_unordered(MAX_PREFETCH_TRANSFERS);
1675            while let Some(result) = transfers.next().await {
1676                result?;
1677            }
1678            pending.extend(
1679                directory
1680                    .directories
1681                    .into_iter()
1682                    .map(|directory| directory.digest),
1683            );
1684        }
1685        Ok(())
1686    }
1687
1688    async fn fetch_remote_blob(
1689        &self,
1690        remote: &RemoteCacheClient,
1691        digest: &CacheDigest,
1692    ) -> Result<PathBuf> {
1693        self.fetch_remote_blob_with_limit(remote, digest, None)
1694            .await
1695    }
1696
1697    async fn fetch_remote_blob_with_limit(
1698        &self,
1699        remote: &RemoteCacheClient,
1700        digest: &CacheDigest,
1701        prefetch_limit: Option<&tokio::sync::Semaphore>,
1702    ) -> Result<PathBuf> {
1703        let lock = self.write_lock(digest);
1704        let _guard = lock.lock().await;
1705        if let Some(path) = self.find_verified_blob(digest)? {
1706            return Ok(path);
1707        }
1708        let _prefetch_permit = match prefetch_limit {
1709            Some(limit) => Some(limit.acquire().await?),
1710            None => None,
1711        };
1712        let _permit = self.remote_transfers.acquire().await?;
1713        let reservation = self.reserve_remote_download(digest.size)?;
1714        self.stats
1715            .remote_blob_requests
1716            .fetch_add(1, Ordering::Relaxed);
1717        let transfer_timer =
1718            AtomicDurationTimer::start(&self.stats.remote_blob_transfer_duration_ns);
1719        let temporary = remote
1720            .get_blob_file(digest, self.remote_staging_dir.as_path())
1721            .await?;
1722        drop(transfer_timer);
1723        let _cas_timer = AtomicDurationTimer::start(&self.stats.local_cas_write_duration_ns);
1724        let path = self.cas.store_verified_file(digest, temporary.path())?;
1725        reservation.commit(digest.size);
1726        self.remember_verified_blob(digest, &path);
1727        self.stats.stores.fetch_add(1, Ordering::Relaxed);
1728        self.stats
1729            .stored_bytes
1730            .fetch_add(digest.size, Ordering::Relaxed);
1731        self.stats
1732            .downloaded_bytes
1733            .fetch_add(digest.size, Ordering::Relaxed);
1734        Ok(path)
1735    }
1736
1737    async fn respond(&self, request: AgentRequest) -> AgentResponse {
1738        let result = match request {
1739            AgentRequest::FindBlob { digest } => self.find_blob(&digest).await,
1740            AgentRequest::FindBlobs { digests } => self.find_blobs(digests).await,
1741            AgentRequest::StoreBlob { digest, source } => self.store_blob(&digest, &source).await,
1742            AgentRequest::FindActionResult { action } => {
1743                self.stats.lookups.fetch_add(1, Ordering::Relaxed);
1744                self.find_action_result(&action).await
1745            }
1746            AgentRequest::RecordActionHit { action, restore } => {
1747                self.record_action_hit(&action, restore)
1748            }
1749            AgentRequest::RecordBypass { kind } => {
1750                *self.stats.bypasses.lock().unwrap().entry(kind).or_insert(0) += 1;
1751                Ok(AgentResponse::BypassRecorded)
1752            }
1753            AgentRequest::RecordUnconsulted => {
1754                self.stats.unconsulted.fetch_add(1, Ordering::Relaxed);
1755                Ok(AgentResponse::UnconsultedRecorded)
1756            }
1757            AgentRequest::RecordCompilerInvocation {
1758                outcome,
1759                crate_name,
1760                duration_ns,
1761            } => self.record_compiler_invocation(&outcome, crate_name.as_deref(), duration_ns),
1762            AgentRequest::RecordActionVerification { matched, restore } => {
1763                self.record_materialization(restore);
1764                self.stats.verifications.fetch_add(1, Ordering::Relaxed);
1765                if !matched {
1766                    self.stats.divergences.fetch_add(1, Ordering::Relaxed);
1767                }
1768                Ok(AgentResponse::ActionVerificationRecorded)
1769            }
1770            AgentRequest::StoreActionResult { result } => self.store_action_result(&result).await,
1771            AgentRequest::FindActionPrediction { task, invocation } => {
1772                self.find_action_prediction(&task, &invocation)
1773            }
1774            AgentRequest::RecordActionPrediction { task, prediction } => {
1775                self.record_action_prediction(&task, prediction)
1776            }
1777            AgentRequest::FindExecutableIdentity {
1778                executable,
1779                environment,
1780            } => self.find_executable_identity(executable, environment),
1781            AgentRequest::StoreExecutableIdentity {
1782                executable,
1783                environment,
1784                stdout,
1785            } => self.store_executable_identity(executable, environment, stdout),
1786            AgentRequest::Hello { .. } => {
1787                Err(eyre::eyre!("hello is only valid as the first request"))
1788            }
1789        };
1790        result.unwrap_or_else(|error| AgentResponse::Error {
1791            message: error.to_string(),
1792        })
1793    }
1794
1795    async fn find_blob(&self, digest: &CacheDigest) -> Result<AgentResponse> {
1796        if let Some(path) = self.find_verified_blob(digest)? {
1797            return Ok(AgentResponse::Blob { path: Some(path) });
1798        }
1799        if !self.remote_mode.reads() {
1800            return Ok(AgentResponse::Blob { path: None });
1801        }
1802        let Some(remote) = &self.remote else {
1803            return Ok(AgentResponse::Blob { path: None });
1804        };
1805        match self.fetch_remote_blob(remote, digest).await {
1806            Ok(path) => Ok(AgentResponse::Blob { path: Some(path) }),
1807            Err(error) => {
1808                warn!(
1809                    "remote cache blob lookup failed for {}: {error}",
1810                    digest.hash
1811                );
1812                Ok(AgentResponse::Blob { path: None })
1813            }
1814        }
1815    }
1816
1817    async fn find_blobs(&self, digests: Vec<CacheDigest>) -> Result<AgentResponse> {
1818        let mut paths = BTreeMap::new();
1819        let mut missing = Vec::new();
1820        for digest in &digests {
1821            match self.find_verified_blob(digest)? {
1822                Some(path) => {
1823                    paths.insert(digest.clone(), path);
1824                }
1825                None => {
1826                    missing.push(digest.clone());
1827                }
1828            }
1829        }
1830
1831        if !missing.is_empty()
1832            && self.remote_mode.reads()
1833            && let Some(remote) = &self.remote
1834        {
1835            paths.extend(self.fetch_remote_blobs(remote, missing, None).await);
1836        }
1837
1838        Ok(AgentResponse::Blobs {
1839            paths: digests
1840                .into_iter()
1841                .map(|digest| paths.get(&digest).cloned())
1842                .collect(),
1843        })
1844    }
1845
1846    async fn store_blob(&self, digest: &CacheDigest, source: &Path) -> Result<AgentResponse> {
1847        let remote = if self.remote_mode.writes() {
1848            self.remote.as_deref()
1849        } else {
1850            None
1851        };
1852        let path = {
1853            let lock = self.write_lock(digest);
1854            let _guard = lock.lock().await;
1855            if let Some(path) = self.find_verified_blob(digest)? {
1856                path
1857            } else {
1858                let path = self.cas.store_file(digest, source)?;
1859                self.remember_verified_blob(digest, &path);
1860                self.stats.stores.fetch_add(1, Ordering::Relaxed);
1861                self.stats
1862                    .stored_bytes
1863                    .fetch_add(digest.size, Ordering::Relaxed);
1864                path
1865            }
1866        };
1867        if let Some(remote) = remote {
1868            let _permit = self.remote_transfers.acquire().await?;
1869            if let Err(error) = remote
1870                .put_blob(&BlobUpload {
1871                    digest: digest.clone(),
1872                    source: BlobSource::Path(path.clone()),
1873                })
1874                .await
1875            {
1876                self.note_remote_failure();
1877                warn!(
1878                    "remote cache blob upload failed for {}: {error}",
1879                    digest.hash
1880                );
1881            } else {
1882                self.stats
1883                    .uploaded_bytes
1884                    .fetch_add(digest.size, Ordering::Relaxed);
1885            }
1886        }
1887        Ok(AgentResponse::Stored { path })
1888    }
1889
1890    fn find_verified_blob(&self, digest: &CacheDigest) -> Result<Option<PathBuf>> {
1891        let remembered = self.verified_blobs.lock().unwrap().get(digest).cloned();
1892        if let Some(path) = remembered {
1893            if digest.matches_file(&path).unwrap_or(false) {
1894                return Ok(Some(path));
1895            }
1896            self.verified_blobs.lock().unwrap().remove(digest);
1897        }
1898        let path = self.cas.find(digest)?;
1899        if let Some(path) = &path {
1900            self.remember_verified_blob(digest, path);
1901        }
1902        Ok(path)
1903    }
1904
1905    fn remember_verified_blob(&self, digest: &CacheDigest, path: &Path) {
1906        self.verified_blobs
1907            .lock()
1908            .unwrap()
1909            .insert(digest.clone(), path.to_path_buf());
1910    }
1911
1912    async fn find_action_result(&self, action: &CacheDigest) -> Result<AgentResponse> {
1913        if let Some(result) = self.actions.find(action)? {
1914            return Ok(AgentResponse::ActionResult {
1915                result: Some(result),
1916            });
1917        }
1918        if !self.remote_mode.reads() {
1919            return Ok(AgentResponse::ActionResult { result: None });
1920        }
1921        let Some(remote) = &self.remote else {
1922            return Ok(AgentResponse::ActionResult { result: None });
1923        };
1924        let lock = self.action_lock(action);
1925        let _guard = lock.lock().await;
1926        if let Some(result) = self.actions.find(action)? {
1927            return Ok(AgentResponse::ActionResult {
1928                result: Some(result),
1929            });
1930        }
1931        if let Some(result) = self
1932            .pending_remote_actions
1933            .lock()
1934            .unwrap()
1935            .get(action)
1936            .cloned()
1937        {
1938            return Ok(AgentResponse::ActionResult {
1939                result: Some(result),
1940            });
1941        }
1942        let _permit = self.remote_transfers.acquire().await?;
1943        match self.get_remote_action_result(remote, action).await {
1944            Ok(Some(result)) => {
1945                self.pending_remote_actions
1946                    .lock()
1947                    .unwrap()
1948                    .insert(action.clone(), result.clone());
1949                Ok(AgentResponse::ActionResult {
1950                    result: Some(result),
1951                })
1952            }
1953            Ok(None) => Ok(AgentResponse::ActionResult { result: None }),
1954            Err(error) => {
1955                self.note_remote_failure();
1956                warn!(
1957                    "remote cache action lookup failed for {}: {error}",
1958                    action.hash
1959                );
1960                Ok(AgentResponse::ActionResult { result: None })
1961            }
1962        }
1963    }
1964
1965    async fn store_action_result(&self, result: &RemoteActionResult) -> Result<AgentResponse> {
1966        let path = self.actions.store(result)?;
1967        if self.remote_mode.writes()
1968            && let Some(remote) = &self.remote
1969        {
1970            let _permit = self.remote_transfers.acquire().await?;
1971            if let Err(error) = remote.put_action_result(result).await {
1972                self.note_remote_failure();
1973                warn!(
1974                    "remote cache action upload failed for {}: {error}",
1975                    result.action.hash
1976                );
1977            }
1978        }
1979        Ok(AgentResponse::ActionStored { path })
1980    }
1981
1982    /// Record that a remote operation failed and the build carried on without it.
1983    fn note_remote_failure(&self) {
1984        self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
1985    }
1986
1987    async fn get_remote_action_result(
1988        &self,
1989        remote: &RemoteCacheClient,
1990        action: &CacheDigest,
1991    ) -> Result<Option<RemoteActionResult>> {
1992        self.stats
1993            .remote_action_lookups
1994            .fetch_add(1, Ordering::Relaxed);
1995        let _timer = AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
1996        remote.get_action_result(action).await
1997    }
1998
1999    fn record_action_hit(
2000        &self,
2001        action: &CacheDigest,
2002        restore: RestoreStats,
2003    ) -> Result<AgentResponse> {
2004        if self.actions.find(action)?.is_none() {
2005            let pending = self.pending_remote_actions.lock().unwrap().remove(action);
2006            if let Some(result) = pending {
2007                self.actions.store(&result)?;
2008            } else {
2009                bail!("cannot record a hit for a missing action result");
2010            }
2011        }
2012        self.record_restore(restore);
2013        self.stats.hits.fetch_add(1, Ordering::Relaxed);
2014        Ok(AgentResponse::ActionHitRecorded)
2015    }
2016
2017    fn record_restore(&self, restore: RestoreStats) {
2018        self.record_materialization(restore);
2019        atomic_saturating_add(
2020            &self.stats.avoided_compiler_duration_ns,
2021            restore.avoided_compiler_duration_ns,
2022        );
2023        atomic_saturating_add(&self.stats.restored_output_files, restore.output_files);
2024        atomic_saturating_add(&self.stats.restored_output_bytes, restore.output_bytes);
2025        atomic_saturating_add(
2026            &self.stats.reflinked_output_files,
2027            restore.reflinked_output_files,
2028        );
2029        atomic_saturating_add(
2030            &self.stats.reflinked_output_bytes,
2031            restore.reflinked_output_bytes,
2032        );
2033        atomic_saturating_add(&self.stats.copied_output_files, restore.copied_output_files);
2034        atomic_saturating_add(&self.stats.copied_output_bytes, restore.copied_output_bytes);
2035    }
2036
2037    fn record_compiler_invocation(
2038        &self,
2039        outcome: &str,
2040        crate_name: Option<&str>,
2041        duration_ns: u64,
2042    ) -> Result<AgentResponse> {
2043        if !matches!(outcome, "miss" | "unconsulted" | "bypass" | "verification") {
2044            bail!("invalid compiler invocation outcome");
2045        }
2046        if let Some(crate_name) = crate_name
2047            && (crate_name.len() > 256 || crate_name.contains(['\0', '\n', '\r']))
2048        {
2049            bail!("invalid compiler crate name");
2050        }
2051        let mut compiler = self.stats.compiler.lock().unwrap();
2052        let stats = compiler.entry(outcome.to_string()).or_default();
2053        stats.invocations = stats.invocations.saturating_add(1);
2054        stats.duration_ns = stats.duration_ns.saturating_add(duration_ns);
2055        drop(compiler);
2056        if outcome != "verification"
2057            && let Some(crate_name) = crate_name.filter(|name| !name.is_empty())
2058        {
2059            let mut slow = self.stats.slow_compilations.lock().unwrap();
2060            let duration = slow.entry(crate_name.to_string()).or_default();
2061            *duration = duration.saturating_add(duration_ns);
2062        }
2063        Ok(AgentResponse::CompilerInvocationRecorded)
2064    }
2065
2066    fn record_materialization(&self, restore: RestoreStats) {
2067        atomic_saturating_add(&self.stats.materialization_duration_ns, restore.duration_ns);
2068    }
2069
2070    fn find_action_prediction(
2071        &self,
2072        task: &str,
2073        invocation: &CacheDigest,
2074    ) -> Result<AgentResponse> {
2075        validate_task_identity(task)?;
2076        invocation.validate()?;
2077        let prediction = self
2078            .task_actions
2079            .lock()
2080            .unwrap()
2081            .get(task)
2082            .and_then(|state| state.predictions.get(invocation))
2083            .cloned();
2084        Ok(AgentResponse::ActionPrediction { prediction })
2085    }
2086
2087    fn record_action_prediction(
2088        &self,
2089        task: &str,
2090        prediction: ActionPrediction,
2091    ) -> Result<AgentResponse> {
2092        validate_task_identity(task)?;
2093        validate_action_prediction(&prediction)?;
2094        let mut tasks = self.task_actions.lock().unwrap();
2095        let state = tasks.entry(task.to_string()).or_default();
2096        if !state.predictions.contains_key(&prediction.invocation)
2097            && state.predictions.len() >= MAX_TASK_ACTION_PREDICTIONS
2098        {
2099            bail!("task action manifest contains too many predictions");
2100        }
2101        state
2102            .predictions
2103            .insert(prediction.invocation.clone(), prediction.clone());
2104        state
2105            .pending_predictions
2106            .insert(prediction.invocation.clone(), prediction);
2107        Ok(AgentResponse::ActionPredictionRecorded)
2108    }
2109
2110    fn executable_identity_key(
2111        &self,
2112        executable: PathBuf,
2113        environment: BTreeMap<String, Option<String>>,
2114    ) -> Result<ExecutableIdentityKey> {
2115        if !environment
2116            .keys()
2117            .all(|name| matches!(name.as_str(), "RUSTUP_HOME" | "RUSTUP_TOOLCHAIN"))
2118        {
2119            bail!("executable identity contains an unsupported environment variable");
2120        }
2121        Ok(ExecutableIdentityKey {
2122            executable,
2123            environment,
2124        })
2125    }
2126
2127    fn find_executable_identity(
2128        &self,
2129        executable: PathBuf,
2130        environment: BTreeMap<String, Option<String>>,
2131    ) -> Result<AgentResponse> {
2132        let key = self.executable_identity_key(executable, environment)?;
2133        let stdout = self
2134            .executable_identities
2135            .lock()
2136            .unwrap()
2137            .get(&key)
2138            .cloned();
2139        Ok(AgentResponse::ExecutableIdentity { stdout })
2140    }
2141
2142    fn store_executable_identity(
2143        &self,
2144        executable: PathBuf,
2145        environment: BTreeMap<String, Option<String>>,
2146        stdout: Vec<u8>,
2147    ) -> Result<AgentResponse> {
2148        if stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE {
2149            bail!("executable identity exceeds {MAX_EXECUTABLE_IDENTITY_SIZE} bytes");
2150        }
2151        let key = self.executable_identity_key(executable, environment)?;
2152        let mut identities = self.executable_identities.lock().unwrap();
2153        let is_new = !identities.contains_key(&key);
2154        let previous_size = identities.get(&key).map_or(0, Vec::len);
2155        if is_new && identities.len() >= MAX_EXECUTABLE_IDENTITIES {
2156            bail!("executable identity cache contains too many entries");
2157        }
2158        let retained_bytes = identities.values().map(Vec::len).sum::<usize>();
2159        if retained_bytes - previous_size + stdout.len() > MAX_EXECUTABLE_IDENTITY_BYTES {
2160            bail!("executable identity cache contains too many bytes");
2161        }
2162        identities.insert(key, stdout.clone());
2163        Ok(AgentResponse::ExecutableIdentity {
2164            stdout: Some(stdout),
2165        })
2166    }
2167
2168    /// Serve newline-delimited protocol requests on an authenticated session stream.
2169    pub async fn handle_connection<S>(&self, stream: S) -> Result<()>
2170    where
2171        S: AsyncRead + AsyncWrite + Unpin,
2172    {
2173        let (reader, mut writer) = tokio::io::split(stream);
2174        let mut reader = BufReader::new(reader);
2175        let hello = read_request(&mut reader)
2176            .await?
2177            .ok_or_else(|| eyre::eyre!("connection closed before the agent handshake"))?;
2178        let request: AgentRequest = serde_json::from_str(&hello)?;
2179        match request {
2180            AgentRequest::Hello {
2181                protocol,
2182                client_version,
2183            } if protocol == AGENT_PROTOCOL_VERSION && client_version == self.version.as_ref() => {}
2184            AgentRequest::Hello { protocol, .. } if protocol != AGENT_PROTOCOL_VERSION => {
2185                send_response(
2186                    &mut writer,
2187                    &AgentResponse::Error {
2188                        message: format!(
2189                            "unsupported agent protocol {protocol}; expected {AGENT_PROTOCOL_VERSION}"
2190                        ),
2191                    },
2192                )
2193                .await?;
2194                return Ok(());
2195            }
2196            AgentRequest::Hello { client_version, .. } => {
2197                send_response(
2198                    &mut writer,
2199                    &AgentResponse::Error {
2200                        message: format!(
2201                            "cache client {client_version} does not match agent {}",
2202                            self.version
2203                        ),
2204                    },
2205                )
2206                .await?;
2207                return Ok(());
2208            }
2209            _ => bail!("the first agent request must be hello"),
2210        }
2211        send_response(
2212            &mut writer,
2213            &AgentResponse::Hello {
2214                protocol: AGENT_PROTOCOL_VERSION,
2215                agent_version: self.version.to_string(),
2216            },
2217        )
2218        .await?;
2219
2220        while let Some(line) = read_request(&mut reader).await? {
2221            let response = match serde_json::from_str(&line) {
2222                Ok(request) => self.respond(request).await,
2223                Err(error) => AgentResponse::Error {
2224                    message: format!("invalid agent request: {error}"),
2225                },
2226            };
2227            send_response(&mut writer, &response).await?;
2228        }
2229        Ok(())
2230    }
2231}
2232
2233/// Read one newline-delimited request, refusing one that grows past the cap.
2234///
2235/// Any process running as this user can open the session socket, so a request
2236/// that never terminates its line must not be able to grow the agent's memory
2237/// without bound.
2238async fn read_request<R>(reader: &mut R) -> Result<Option<String>>
2239where
2240    R: AsyncBufRead + Unpin,
2241{
2242    let mut line = Vec::new();
2243    loop {
2244        let available = reader.fill_buf().await?;
2245        if available.is_empty() {
2246            break;
2247        }
2248        let (consumed, complete) = match available.iter().position(|byte| *byte == b'\n') {
2249            Some(index) => (index, true),
2250            None => (available.len(), false),
2251        };
2252        if line.len() + consumed > MAX_REQUEST_BYTES {
2253            bail!("agent request exceeded {MAX_REQUEST_BYTES} bytes");
2254        }
2255        line.extend_from_slice(&available[..consumed]);
2256        // The newline itself is consumed but never kept.
2257        reader.consume(consumed + usize::from(complete));
2258        if complete {
2259            return Ok(Some(String::from_utf8(line)?));
2260        }
2261    }
2262    if line.is_empty() {
2263        Ok(None)
2264    } else {
2265        Ok(Some(String::from_utf8(line)?))
2266    }
2267}
2268
2269/// Whether `task` is a well-formed task action identity.
2270///
2271/// Identities name files and directories in the store, so anything that reads
2272/// the store back has to be able to tell an identity from whatever else a user
2273/// left lying there.
2274pub fn is_task_identity(task: &str) -> bool {
2275    task.len() == 64
2276        && task
2277            .bytes()
2278            .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
2279}
2280
2281fn validate_task_identity(task: &str) -> Result<()> {
2282    if !is_task_identity(task) {
2283        bail!("invalid task action identity");
2284    }
2285    Ok(())
2286}
2287
2288/// Where a store keeps its task prediction manifests.
2289fn task_manifest_dir(store: &Path) -> PathBuf {
2290    store.join("task-manifests").join("v1")
2291}
2292
2293/// The action digests a task's prediction manifest recorded.
2294///
2295/// Read straight off disk rather than through an agent, because a collector
2296/// needs the action set of tasks no session is running. A manifest that is
2297/// missing or no longer parseable yields no actions rather than an error: this
2298/// is a prediction index, so the worst a thin answer costs is a cold prefetch,
2299/// or an object collected earlier than it deserved.
2300pub fn task_manifest_actions(store: &Path, task: &str) -> Result<Vec<CacheDigest>> {
2301    validate_task_identity(task)?;
2302    let path = task_manifest_dir(store).join(format!("{task}.json"));
2303    let bytes = match fs::read(&path) {
2304        Ok(bytes) => bytes,
2305        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
2306        Err(error) => {
2307            return Err(error).wrap_err_with(|| format!("failed to read {}", path.display()));
2308        }
2309    };
2310    let Ok(manifest) = serde_json::from_slice::<TaskActionManifest>(&bytes) else {
2311        return Ok(Vec::new());
2312    };
2313    if validate_task_manifest(&manifest, task).is_err() {
2314        return Ok(Vec::new());
2315    }
2316    Ok(manifest
2317        .predictions
2318        .into_iter()
2319        .map(|prediction| prediction.action)
2320        .collect())
2321}
2322
2323fn validate_action_prediction(prediction: &ActionPrediction) -> Result<()> {
2324    if prediction.validate() {
2325        Ok(())
2326    } else {
2327        bail!("invalid action prediction")
2328    }
2329}
2330
2331fn validate_task_manifest(manifest: &TaskActionManifest, task: &str) -> Result<()> {
2332    if manifest.task == task && manifest.validate() {
2333        Ok(())
2334    } else {
2335        bail!("invalid task action manifest")
2336    }
2337}
2338
2339fn merge_task_manifests(
2340    task: &str,
2341    base: Option<TaskActionManifest>,
2342    update: TaskActionManifest,
2343) -> Result<TaskActionManifest> {
2344    validate_task_manifest(&update, task)?;
2345    let mut predictions = BTreeMap::new();
2346    if let Some(base) = base {
2347        validate_task_manifest(&base, task)?;
2348        predictions.extend(
2349            base.predictions
2350                .into_iter()
2351                .map(|prediction| (prediction.invocation.clone(), prediction)),
2352        );
2353    }
2354    predictions.extend(
2355        update
2356            .predictions
2357            .into_iter()
2358            .map(|prediction| (prediction.invocation.clone(), prediction)),
2359    );
2360    let manifest = TaskActionManifest {
2361        version: TASK_ACTION_MANIFEST_VERSION,
2362        task: task.to_owned(),
2363        predictions: predictions.into_values().collect(),
2364    };
2365    validate_task_manifest(&manifest, task)?;
2366    Ok(manifest)
2367}
2368
2369fn merge_remote_task_manifest(
2370    task: &str,
2371    remote: TaskActionManifest,
2372    local: TaskActionManifest,
2373) -> (TaskActionManifest, bool) {
2374    match merge_task_manifests(task, Some(remote), local.clone()) {
2375        Ok(manifest) => (manifest, true),
2376        Err(error) => {
2377            warn!("remote task action manifest merge failed for {task}: {error}");
2378            (local, false)
2379        }
2380    }
2381}
2382
2383async fn send_response(
2384    writer: &mut (impl AsyncWrite + Unpin),
2385    response: &AgentResponse,
2386) -> Result<()> {
2387    let mut encoded = serde_json::to_vec(response)?;
2388    encoded.push(b'\n');
2389    writer.write_all(&encoded).await?;
2390    writer.flush().await?;
2391    Ok(())
2392}
2393
2394#[cfg(test)]
2395#[path = "agent_tests.rs"]
2396mod tests;