Skip to main content

mbx_cache_core/
agent.rs

1use crate::uploads::{ConnectionUploads, UploadQueue, UploadSink};
2use crate::{
3    ActionPrediction, BlobPackLimits, CacheDigest, CacheDirectory, LocalActionCache, LocalCas,
4    MAX_STAGED_BLOB_PACK_BYTES, MAX_STAGED_BLOB_PACK_ITEMS, ManifestPutOutcome, RemoteActionResult,
5    RemoteCacheClient, RemoteCacheMode, RustcMetadata, TaskActionManifest, blob_pack_chunk,
6    canonical_json,
7};
8use eyre::{Context, Result, bail};
9use futures_util::{FutureExt, StreamExt, future::BoxFuture, stream};
10use log::warn;
11use serde::{Deserialize, Serialize};
12use std::collections::{BTreeMap, BTreeSet};
13use std::fs;
14use std::path::{Path, PathBuf};
15use std::sync::atomic::{AtomicU64, Ordering};
16use std::sync::{Arc, Mutex, Weak};
17use std::time::{Duration, Instant, SystemTime};
18use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncRead, AsyncWrite, AsyncWriteExt, BufReader};
19
20const MAX_EXECUTABLE_IDENTITIES: usize = 64;
21const MAX_EXECUTABLE_IDENTITY_SIZE: usize = 64 * 1024;
22const MAX_EXECUTABLE_IDENTITY_BYTES: usize = 256 * 1024;
23const TASK_ACTION_MANIFEST_VERSION: u8 = 1;
24const MAX_TASK_ACTION_PREDICTIONS: usize = 16 * 1024;
25/// Longest shim diagnostic the agent accepts, in bytes.
26const MAX_WARNING_BYTES: usize = 4 * 1024;
27/// Most distinct shim diagnostics one session surfaces before going quiet.
28///
29/// A failure mode that repeats across a build tends to repeat with the same
30/// message, which deduplication already collapses; the cap only guards
31/// against a message that embeds something unique per compilation.
32const MAX_WARNINGS: usize = 128;
33/// Most file identities one digest lookup or record may carry.
34const MAX_FILE_DIGEST_BATCH: usize = 16 * 1024;
35/// Most recorded file digests one session retains across every scope.
36///
37/// Roughly two orders of magnitude above the largest workspace measured; the
38/// cap exists so a pathological build bounds the agent instead of growing it.
39const MAX_FILE_DIGEST_ENTRIES: usize = 1024 * 1024;
40const MAX_REMOTE_TRANSFERS: usize = 64;
41const MAX_PREFETCH_TRANSFERS: usize = 48;
42const MAX_PREFETCH_ACTION_BATCH: usize = 256;
43/// Batched action lookups issued at once, well inside the transfer budget: each
44/// asks about hundreds of actions, so a handful covers a large workspace.
45const MAX_PREFETCH_BATCH_LOOKUPS: usize = 4;
46const PREFETCH_ACTION_BATCH_DELAY: Duration = Duration::from_millis(5);
47const MAX_PREFETCH_DIRECTORY_OBJECTS: usize = 100_000;
48const MAX_PREFETCH_OBJECTS_PER_WAVE: usize = 100_000;
49const DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES: u64 = 5 * 1024 * 1024 * 1024;
50
51/// Remote action-cache access owned by one task session.
52pub struct AgentRemoteCache {
53    /// Remote protocol client used by the agent.
54    pub client: RemoteCacheClient,
55    /// Permitted remote read/write operations.
56    pub mode: RemoteCacheMode,
57    /// Directory used for verified downloads before CAS ingestion.
58    pub staging_dir: PathBuf,
59}
60
61/// Wire protocol version used between an in-process cache agent and its shims.
62pub const AGENT_PROTOCOL_VERSION: u8 = 5;
63/// Largest single protocol request the agent will read.
64///
65/// Requests are small JSON objects; the largest legitimate ones carry an output
66/// tree or a batch of digests, which stay far below this.
67const MAX_REQUEST_BYTES: usize = 16 * 1024 * 1024;
68
69/// A request accepted by the task-scoped cache agent.
70#[derive(Debug, Clone, Serialize, Deserialize)]
71#[serde(tag = "type", rename_all = "snake_case")]
72pub enum AgentRequest {
73    /// Negotiate protocol and application versions.
74    Hello {
75        /// Agent protocol version understood by the caller.
76        protocol: u8,
77        /// Human-readable mbx client version.
78        client_version: String,
79    },
80    /// Begin a prediction-manifest run for a stable Cargo or task identity.
81    BeginTask {
82        /// Stable 64-character lowercase hexadecimal task identity.
83        task: String,
84    },
85    /// Commit predictions collected by an earlier [`Self::BeginTask`].
86    CommitTask {
87        /// Opaque run identifier returned by the agent.
88        run: String,
89    },
90    /// Resolve a blob to a session-verified local CAS path.
91    FindBlob {
92        /// Blob to resolve.
93        digest: CacheDigest,
94    },
95    /// Resolve blobs to session-verified local CAS paths.
96    FindBlobs {
97        /// Blobs to resolve, preserving request order in the response.
98        digests: Vec<CacheDigest>,
99    },
100    /// Import a file into the local content-addressed store.
101    StoreBlob {
102        /// Digest the source must match.
103        digest: CacheDigest,
104        /// File to verify and import.
105        source: PathBuf,
106    },
107    /// Look up an action-result record.
108    FindActionResult {
109        /// Action digest to resolve.
110        action: CacheDigest,
111    },
112    /// Account for a successfully restored cache hit.
113    RecordActionHit {
114        /// Action that supplied the outputs.
115        action: CacheDigest,
116        /// Restoration work performed by the adapter.
117        restore: RestoreStats,
118        /// Compiler crate name, when the invocation supplied one.
119        #[serde(default)]
120        crate_name: Option<String>,
121    },
122    /// A compilation the adapter declined to cache, grouped by reason.
123    RecordBypass {
124        /// Stable, low-cardinality bypass-reason name.
125        kind: String,
126    },
127    /// A compilation the adapter could not look up, having no key to look up
128    /// with. Distinct from a bypass: these are cached once compiled.
129    RecordUnconsulted,
130    /// Account for one real compiler invocation performed by the adapter.
131    RecordCompilerInvocation {
132        /// Stable outcome category such as `miss`, `unconsulted`, `bypass`, or
133        /// `incremental` for a compilation the adapter deliberately ran with
134        /// incremental state instead of publishing.
135        outcome: String,
136        /// Compiler crate name, when the invocation supplied one.
137        crate_name: Option<String>,
138        /// Wall time spent running the compiler.
139        duration_ns: u64,
140    },
141    /// Account for a cache hit that was rebuilt for correctness verification.
142    RecordActionVerification {
143        /// Whether rebuilt and cached outputs matched.
144        matched: bool,
145        /// Restoration work performed before rebuilding.
146        restore: RestoreStats,
147    },
148    /// Store an action-result record locally and enqueue remote publication.
149    StoreActionResult {
150        /// Action-result record to store.
151        result: RemoteActionResult,
152    },
153    /// Find an earlier input prediction for a task and invocation.
154    FindActionPrediction {
155        /// Stable task identity.
156        task: String,
157        /// Digest of the compiler invocation without discovered inputs.
158        invocation: CacheDigest,
159    },
160    /// Record an input prediction after a successful compilation.
161    RecordActionPrediction {
162        /// Stable task identity.
163        task: String,
164        /// Adapter-owned prediction record.
165        prediction: ActionPrediction,
166    },
167    /// Find cached identity output for an executable and environment.
168    FindExecutableIdentity {
169        /// Executable whose identity command would run.
170        executable: PathBuf,
171        /// Environment variables affecting identity output.
172        environment: BTreeMap<String, Option<String>>,
173    },
174    /// Cache identity output for an executable and environment.
175    StoreExecutableIdentity {
176        /// Executable whose identity command ran.
177        executable: PathBuf,
178        /// Environment variables affecting identity output.
179        environment: BTreeMap<String, Option<String>>,
180        /// Captured identity-command standard output.
181        stdout: Vec<u8>,
182    },
183    /// Surface a shim diagnostic through the session that owns the build.
184    ///
185    /// A shim must not print diagnostics itself: its stderr belongs to the
186    /// compiler it stands in for, and build scripts read that stream as part
187    /// of the compiler's answer -- cc-rs treats any stderr output from a flag
188    /// probe as "unsupported", which changes the flags of every compilation
189    /// that follows and, with them, every action key the build produces.
190    ///
191    /// Appended rather than grouped with the other `Record*` requests it
192    /// belongs beside: these variants carry no `repr`, so inserting one moves
193    /// the discriminant of every variant after it, and a break nobody asked
194    /// for is worth less than the grouping. New variants go here.
195    RecordWarning {
196        /// Human-readable single-line diagnostic.
197        message: String,
198    },
199    /// Find session-recorded digests for files with these identities.
200    FindFileDigests {
201        /// What the recorded digests may stand in for.
202        scope: FileDigestScope,
203        /// File identities to resolve, preserving request order in the
204        /// response.
205        files: Vec<FileIdentity>,
206    },
207    /// Record digests of files a shim hashed or wrote this session.
208    RecordFileDigests {
209        /// What the recorded digests may stand in for.
210        scope: FileDigestScope,
211        /// Hashed files and the identities their digests describe.
212        entries: Vec<RecordedFileDigest>,
213    },
214}
215
216/// Local output restoration work performed by one action-cache adapter hit.
217#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
218#[serde(deny_unknown_fields)]
219pub struct RestoreStats {
220    /// Cumulative time spent materializing and validating output files.
221    pub duration_ns: u64,
222    /// Compiler wall time recorded when this action was originally produced.
223    /// Zero means no timing hint was available.
224    pub avoided_compiler_duration_ns: u64,
225    /// Number of compiler output files restored.
226    pub output_files: u64,
227    /// Declared size of compiler output files restored.
228    pub output_bytes: u64,
229    /// Number of restored output files that share data blocks with the CAS.
230    pub reflinked_output_files: u64,
231    /// Declared size of restored outputs that share data blocks with the CAS.
232    pub reflinked_output_bytes: u64,
233    /// Number of restored output files that required a byte-for-byte copy.
234    pub copied_output_files: u64,
235    /// Declared size of restored outputs that required a byte-for-byte copy.
236    pub copied_output_bytes: u64,
237    /// Number of output files already in place with the cached contents, kept
238    /// rather than rewritten.
239    pub reused_output_files: u64,
240    /// Declared size of outputs kept in place rather than rewritten.
241    pub reused_output_bytes: u64,
242}
243
244/// One accounted cache decision, as it happens.
245///
246/// The agent already folds every one of these into [`AgentStats`]; an observer
247/// sees the same decisions individually, before that summing loses the crate
248/// they belong to. Delivered synchronously from the request handler, so an
249/// observer that blocks slows the build it is watching.
250#[derive(Debug, Clone)]
251#[non_exhaustive]
252pub enum AgentEvent {
253    /// An action's outputs were restored from cache.
254    ActionHit {
255        /// Compiler crate name, when the invocation supplied one.
256        crate_name: Option<String>,
257        /// Restoration work performed by the adapter.
258        restore: RestoreStats,
259    },
260    /// A compilation the adapter declined to cache.
261    Bypass {
262        /// Stable, low-cardinality bypass-reason name.
263        kind: String,
264    },
265    /// A compilation no lookup was possible for.
266    Unconsulted,
267    /// A real compiler invocation ran.
268    CompilerInvocation {
269        /// Stable outcome category such as `miss`, `unconsulted`, or `bypass`.
270        outcome: String,
271        /// Compiler crate name, when the invocation supplied one.
272        crate_name: Option<String>,
273        /// Wall time spent running the compiler.
274        duration_ns: u64,
275    },
276    /// A hit was rebuilt to verify it.
277    Verification {
278        /// Whether rebuilt and cached outputs matched.
279        matched: bool,
280        /// Restoration work performed before rebuilding.
281        restore: RestoreStats,
282    },
283    /// A shim reported a diagnostic for the session to surface.
284    Warning {
285        /// Human-readable single-line diagnostic.
286        message: String,
287    },
288}
289
290/// A sink for [`AgentEvent`]s observed during one session.
291pub trait AgentEventObserver: Send + Sync {
292    /// Handle one event. Must not panic, and should not block.
293    fn event(&self, event: AgentEvent);
294}
295
296/// A response returned by the task-scoped cache agent.
297#[derive(Debug, Clone, Serialize, Deserialize)]
298#[serde(tag = "type", rename_all = "snake_case")]
299pub enum AgentResponse {
300    /// Successful protocol negotiation.
301    Hello {
302        /// Agent protocol version.
303        protocol: u8,
304        /// Human-readable agent version.
305        agent_version: String,
306    },
307    /// A task prediction run was begun.
308    TaskBegun {
309        /// Opaque run identifier to pass to compiler shims and commit later.
310        run: String,
311    },
312    /// A task prediction run was committed.
313    TaskCommitted,
314    /// A local CAS path already verified against the requested digest.
315    Blob {
316        /// Verified local path, or `None` on a cache miss.
317        path: Option<PathBuf>,
318    },
319    /// Local CAS paths already verified against the requested digests.
320    Blobs {
321        /// Verified local paths or misses, in request order.
322        paths: Vec<Option<PathBuf>>,
323    },
324    /// A blob was stored locally.
325    Stored {
326        /// Path of the stored object in the local CAS.
327        path: PathBuf,
328    },
329    /// Result of an action lookup.
330    ActionResult {
331        /// Validated action result, or `None` on a cache miss.
332        result: Option<RemoteActionResult>,
333    },
334    /// Hit statistics were updated.
335    ActionHitRecorded,
336    /// Verification statistics were updated.
337    ActionVerificationRecorded,
338    /// Bypass statistics were updated.
339    BypassRecorded,
340    /// Unconsulted-compilation statistics were updated.
341    UnconsultedRecorded,
342    /// Compiler invocation accounting was recorded.
343    CompilerInvocationRecorded,
344    /// An action result was stored.
345    ActionStored {
346        /// Path of the stored local action-result record.
347        path: PathBuf,
348    },
349    /// Result of an input-prediction lookup.
350    ActionPrediction {
351        /// Matching prediction, or `None` when none is known.
352        prediction: Option<ActionPrediction>,
353    },
354    /// An input prediction was recorded.
355    ActionPredictionRecorded,
356    /// Result of an executable-identity lookup.
357    ExecutableIdentity {
358        /// Captured output, or `None` when no identity is cached.
359        stdout: Option<Vec<u8>>,
360    },
361    /// The request failed without terminating the agent connection.
362    Error {
363        /// Human-readable failure description.
364        message: String,
365    },
366    /// A shim diagnostic was accepted for the session to surface.
367    ///
368    /// Sits past `Error` for the reason [`AgentRequest::RecordWarning`] sits
369    /// last: anywhere earlier renumbers the variants below it. New variants go
370    /// here.
371    WarningRecorded,
372    /// Digests recorded earlier for the requested file identities.
373    FileDigests {
374        /// Recorded digests or misses, in request order.
375        digests: Vec<Option<CacheDigest>>,
376    },
377    /// File digests were recorded.
378    FileDigestsRecorded,
379}
380
381/// Aggregate cache activity for one task session.
382///
383/// The agent produces these; nothing outside this crate has cause to build one.
384/// Saying so keeps a new counter from being a breaking change, which is what
385/// this type exists to accumulate -- reach for [`AgentStats::default`] and
386/// assign the fields a test needs.
387#[derive(Debug, Clone, Default, PartialEq, Eq)]
388#[non_exhaustive]
389pub struct AgentStats {
390    /// End-to-end lifetime of the task-scoped cache session.
391    pub session_duration_ns: u64,
392    /// Number of action-result lookups.
393    pub lookups: u64,
394    /// Compilations no action-result lookup was possible for, because no usable
395    /// action key was available.
396    ///
397    /// Counted separately from a miss, which is a lookup that found nothing.
398    /// Both compile, but only a miss says a lookup happened.
399    pub unconsulted: u64,
400    /// Number of lookups that found a valid local action result.
401    pub hits: u64,
402    /// Number of newly stored content-addressed objects.
403    pub stores: u64,
404    /// Total size of newly stored objects.
405    pub stored_bytes: u64,
406    /// Number of cache hits compiled again for qualification.
407    pub verifications: u64,
408    /// Number of qualification builds that diverged from the cached result.
409    pub divergences: u64,
410    /// CAS payload bytes downloaded from the remote cache.
411    pub downloaded_bytes: u64,
412    /// CAS payload bytes uploaded to the remote cache.
413    pub uploaded_bytes: u64,
414    /// Objects published to the remote cache after the build asked for them.
415    ///
416    /// A store request returns once the object is in the local CAS, so the
417    /// upload it implies happens off the build's critical path. This counts the
418    /// blobs and action results that publication actually completed for.
419    pub background_uploads: u64,
420    /// Queued uploads that did not publish, having been reported and recovered
421    /// from.
422    pub background_upload_failures: u64,
423    /// Framed requests that published several blobs at once.
424    pub remote_blob_pack_uploads: u64,
425    /// Blobs published through those framed requests.
426    pub remote_blob_pack_upload_blobs: u64,
427    /// Time the session spent waiting for queued uploads once the build ended.
428    pub upload_drain_duration_ns: u64,
429    /// Complete actions staged before an adapter requested them.
430    pub prefetched_actions: u64,
431    /// Predictions carried by the task manifest this session loaded.
432    ///
433    /// Zero means no earlier build left a manifest behind -- a genuinely cold
434    /// start. Together with `lookups`, this is what tells a session that loaded
435    /// hundreds of predictions and matched none of them apart from one that had
436    /// nothing to match against; the first means the invocations changed (a
437    /// compiler update does this to every one of them at once), the second that
438    /// the store was empty.
439    pub predictions_loaded: u64,
440    /// Compilations that were not cacheable, counted by reason.
441    pub bypasses: BTreeMap<String, u64>,
442    /// Estimated compiler time avoided by restored action hits.
443    pub avoided_compiler_duration_ns: u64,
444    /// Real compiler work performed in this session, grouped by outcome.
445    pub compiler: BTreeMap<String, CompilerStats>,
446    /// Cumulative real compiler time by crate name.
447    pub slow_compilations: BTreeMap<String, u64>,
448    /// Remote cache operations that failed and were degraded to a local result.
449    ///
450    /// A remote cache that cannot be reached, or that answers in a way this
451    /// client refuses, costs hit rate rather than correctness, so every one of
452    /// these is recovered from rather than raised. Counting them is what keeps a
453    /// remote that is failing every request from reading as one that merely had
454    /// nothing to offer.
455    pub remote_failures: u64,
456    /// Number of task manifest requests made to the remote cache.
457    pub remote_manifest_lookups: u64,
458    /// Cumulative time spent requesting remote task manifests.
459    pub remote_manifest_lookup_duration_ns: u64,
460    /// Number of action-result requests made to the remote cache.
461    pub remote_action_lookups: u64,
462    /// Cumulative time spent requesting remote action results.
463    pub remote_action_lookup_duration_ns: u64,
464    /// Number of blob requests made to the remote cache.
465    pub remote_blob_requests: u64,
466    /// Number of packed blob requests made to the remote cache.
467    pub remote_blob_pack_requests: u64,
468    /// Number of verified blobs received through packed responses.
469    pub remote_blob_pack_blobs: u64,
470    /// Cumulative time spent downloading and verifying remote blobs.
471    pub remote_blob_transfer_duration_ns: u64,
472    /// Cumulative time spent ingesting downloaded blobs into the local CAS.
473    pub local_cas_write_duration_ns: u64,
474    /// Number of speculative prefetch runs started for task manifests.
475    pub prefetch_runs: u64,
476    /// Cumulative wall time of speculative task-manifest prefetch runs.
477    pub prefetch_duration_ns: u64,
478    /// Cumulative time spent staging or materializing and validating cached outputs.
479    pub materialization_duration_ns: u64,
480    /// Number of compiler output files restored from action hits.
481    pub restored_output_files: u64,
482    /// Declared size of compiler output files restored from action hits.
483    pub restored_output_bytes: u64,
484    /// Number of restored output files materialized with filesystem reflinks.
485    pub reflinked_output_files: u64,
486    /// Declared size of outputs materialized with filesystem reflinks.
487    pub reflinked_output_bytes: u64,
488    /// Number of restored output files materialized by copying their bytes.
489    pub copied_output_files: u64,
490    /// Declared size of outputs materialized by copying their bytes.
491    pub copied_output_bytes: u64,
492    /// Number of output files kept in place, already holding the cached bytes.
493    pub reused_output_files: u64,
494    /// Declared size of outputs kept in place.
495    pub reused_output_bytes: u64,
496}
497
498/// Count and cumulative wall time for one compiler-invocation outcome.
499///
500/// Non-exhaustive for the same reason as [`AgentStats`], which holds these:
501/// leaving the outer bag open does not help if describing an outcome in more
502/// detail still breaks the type inside it.
503#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
504#[non_exhaustive]
505pub struct CompilerStats {
506    /// Number of compiler invocations observed.
507    pub invocations: u64,
508    /// Cumulative wall time spent in those invocations.
509    pub duration_ns: u64,
510}
511
512impl CompilerStats {
513    /// The counts observed for one outcome.
514    pub fn new(invocations: u64, duration_ns: u64) -> Self {
515        Self {
516            invocations,
517            duration_ns,
518        }
519    }
520}
521
522#[derive(Default)]
523struct AtomicAgentStats {
524    lookups: AtomicU64,
525    unconsulted: AtomicU64,
526    hits: AtomicU64,
527    stores: AtomicU64,
528    stored_bytes: AtomicU64,
529    verifications: AtomicU64,
530    divergences: AtomicU64,
531    downloaded_bytes: AtomicU64,
532    uploaded_bytes: AtomicU64,
533    background_uploads: AtomicU64,
534    background_upload_failures: AtomicU64,
535    remote_blob_pack_uploads: AtomicU64,
536    remote_blob_pack_upload_blobs: AtomicU64,
537    upload_drain_duration_ns: AtomicU64,
538    prefetched_actions: AtomicU64,
539    predictions_loaded: AtomicU64,
540    remote_failures: AtomicU64,
541    remote_manifest_lookups: AtomicU64,
542    remote_manifest_lookup_duration_ns: AtomicU64,
543    remote_action_lookups: AtomicU64,
544    remote_action_lookup_duration_ns: AtomicU64,
545    remote_blob_requests: AtomicU64,
546    remote_blob_pack_requests: AtomicU64,
547    remote_blob_pack_blobs: AtomicU64,
548    remote_blob_transfer_duration_ns: AtomicU64,
549    local_cas_write_duration_ns: AtomicU64,
550    prefetch_runs: AtomicU64,
551    prefetch_duration_ns: AtomicU64,
552    materialization_duration_ns: AtomicU64,
553    bypasses: Mutex<BTreeMap<String, u64>>,
554    avoided_compiler_duration_ns: AtomicU64,
555    compiler: Mutex<BTreeMap<String, CompilerStats>>,
556    slow_compilations: Mutex<BTreeMap<String, u64>>,
557    restored_output_files: AtomicU64,
558    restored_output_bytes: AtomicU64,
559    reflinked_output_files: AtomicU64,
560    reflinked_output_bytes: AtomicU64,
561    copied_output_files: AtomicU64,
562    copied_output_bytes: AtomicU64,
563    reused_output_files: AtomicU64,
564    reused_output_bytes: AtomicU64,
565}
566
567struct AtomicDurationTimer<'a> {
568    started: Instant,
569    target: &'a AtomicU64,
570}
571
572impl<'a> AtomicDurationTimer<'a> {
573    fn start(target: &'a AtomicU64) -> Self {
574        Self {
575            started: Instant::now(),
576            target,
577        }
578    }
579}
580
581impl Drop for AtomicDurationTimer<'_> {
582    fn drop(&mut self) {
583        atomic_saturating_add(self.target, duration_ns(self.started));
584    }
585}
586
587fn duration_ns(started: Instant) -> u64 {
588    started.elapsed().as_nanos().try_into().unwrap_or(u64::MAX)
589}
590
591fn validate_crate_name(crate_name: Option<&str>) -> Result<()> {
592    if let Some(crate_name) = crate_name
593        && (crate_name.len() > 256 || crate_name.contains(['\0', '\n', '\r']))
594    {
595        bail!("invalid compiler crate name");
596    }
597    Ok(())
598}
599
600fn atomic_saturating_add(target: &AtomicU64, value: u64) {
601    let _ = target.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| {
602        Some(current.saturating_add(value))
603    });
604}
605
606fn queue_prefetch_digest(
607    verified: &BTreeMap<CacheDigest, PathBuf>,
608    pending: &mut BTreeMap<CacheDigest, ()>,
609    digest: CacheDigest,
610) {
611    if verified.contains_key(&digest) || pending.contains_key(&digest) {
612        return;
613    }
614    pending.insert(digest, ());
615}
616
617fn queue_prefetch_directory(
618    seen: &BTreeMap<CacheDigest, ()>,
619    pending: &mut BTreeMap<CacheDigest, ()>,
620    digest: CacheDigest,
621    limit: usize,
622) -> bool {
623    if seen.contains_key(&digest) || pending.contains_key(&digest) {
624        return true;
625    }
626    if seen.len().saturating_add(pending.len()) >= limit {
627        return false;
628    }
629    pending.insert(digest, ());
630    true
631}
632
633/// Shared state for an agent hosted by the process that owns a build session.
634///
635/// Transport listeners deliberately live in the embedder so the session
636/// lifecycle owns them. This type only contains ecosystem-independent CAS and
637/// protocol logic.
638#[derive(Clone)]
639pub struct CacheAgent {
640    cas: LocalCas,
641    actions: LocalActionCache,
642    verified_blobs: Arc<Mutex<BTreeMap<CacheDigest, VerifiedBlob>>>,
643    version: Arc<str>,
644    write_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
645    action_locks: Arc<Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>>,
646    stats: Arc<AtomicAgentStats>,
647    observer: Option<Arc<dyn AgentEventObserver>>,
648    executable_identities: Arc<Mutex<BTreeMap<ExecutableIdentityKey, Vec<u8>>>>,
649    manifest_dir: Arc<PathBuf>,
650    task_actions: Arc<Mutex<BTreeMap<String, TaskActionState>>>,
651    next_task_run: Arc<AtomicU64>,
652    manifest_write_lock: Arc<Mutex<()>>,
653    remote: Option<Arc<RemoteCacheClient>>,
654    remote_mode: RemoteCacheMode,
655    remote_staging_dir: Arc<PathBuf>,
656    remote_download_limit: u64,
657    remote_download_bytes: Arc<AtomicU64>,
658    pending_remote_actions: Arc<Mutex<BTreeMap<CacheDigest, RemoteActionResult>>>,
659    remote_transfers: Arc<tokio::sync::Semaphore>,
660    prefetch_transfers: Arc<tokio::sync::Semaphore>,
661    prefetch_tasks: Arc<Mutex<Vec<tokio::task::JoinHandle<()>>>>,
662    /// Distinct shim diagnostics already surfaced, so a warning that fires
663    /// once per compilation is printed once per session.
664    warnings: Arc<Mutex<BTreeSet<String>>>,
665    /// Digests of files shims hashed or wrote this session, keyed by scope and
666    /// path, each entry standing while its recorded identity matches the disk.
667    file_digests: Arc<Mutex<BTreeMap<(FileDigestScope, PathBuf), RecordedFileDigest>>>,
668    /// Deferred remote publication, present only when the session may write.
669    uploads: Option<UploadQueue>,
670}
671
672/// Records background upload activity against a session's statistics.
673struct AgentUploadSink {
674    stats: Arc<AtomicAgentStats>,
675}
676
677impl UploadSink for AgentUploadSink {
678    fn record_blob_uploaded(&self, bytes: u64) {
679        self.stats
680            .background_uploads
681            .fetch_add(1, Ordering::Relaxed);
682        self.stats
683            .uploaded_bytes
684            .fetch_add(bytes, Ordering::Relaxed);
685    }
686
687    fn record_action_uploaded(&self) {
688        self.stats
689            .background_uploads
690            .fetch_add(1, Ordering::Relaxed);
691    }
692
693    fn record_blob_pack_uploaded(&self, blobs: u64) {
694        self.stats
695            .remote_blob_pack_uploads
696            .fetch_add(1, Ordering::Relaxed);
697        self.stats
698            .remote_blob_pack_upload_blobs
699            .fetch_add(blobs, Ordering::Relaxed);
700    }
701
702    fn record_upload_failure(&self) {
703        self.stats
704            .background_upload_failures
705            .fetch_add(1, Ordering::Relaxed);
706        self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
707    }
708}
709
710/// The on-disk identity a recorded file digest describes.
711///
712/// The same trade [`VerifiedBlob`] makes for CAS reads, offered to shims for
713/// the files they hash: an overwrite moves the modification time and a
714/// truncation changes the length, so a digest recorded against both stands
715/// until either does. Where the platform reports a metadata-change time the
716/// identity carries that too, and it is the part a writer cannot restore: a
717/// rewrite that puts the modification time back still moves the change time,
718/// so only filesystems without one fall back to the freshness model the
719/// surrounding build tool already lives on.
720#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
721#[serde(deny_unknown_fields)]
722pub struct FileIdentity {
723    /// Absolute path of the file.
724    pub path: PathBuf,
725    /// Length of the file in bytes.
726    pub len: u64,
727    /// Modification time of the file.
728    pub modified: SystemTime,
729    /// Platform metadata-change token, where one exists.
730    pub changed: Option<(i64, i64)>,
731}
732
733impl FileIdentity {
734    /// Describe a file from metadata already in hand, or nothing when the
735    /// filesystem reports no modification time to compare against later.
736    pub fn describe(path: &Path, metadata: &std::fs::Metadata) -> Option<Self> {
737        Some(Self {
738            path: path.to_path_buf(),
739            len: metadata.len(),
740            modified: metadata.modified().ok()?,
741            changed: change_token(metadata),
742        })
743    }
744}
745
746/// The metadata-change time as an opaque token, where the platform has one.
747#[cfg(unix)]
748fn change_token(metadata: &std::fs::Metadata) -> Option<(i64, i64)> {
749    use std::os::unix::fs::MetadataExt;
750    Some((metadata.ctime(), metadata.ctime_nsec()))
751}
752
753/// Windows reports creation rather than metadata-change time, which a rewrite
754/// does not move, so no token is better than a misleading one.
755#[cfg(not(unix))]
756fn change_token(_metadata: &std::fs::Metadata) -> Option<(i64, i64)> {
757    None
758}
759
760/// A file digest recorded against the identity it was read under.
761#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
762#[serde(deny_unknown_fields)]
763pub struct RecordedFileDigest {
764    /// Identity of the file when its contents were hashed.
765    pub file: FileIdentity,
766    /// Digest of those contents.
767    pub digest: CacheDigest,
768}
769
770/// What a recorded file digest may stand in for.
771///
772/// Adapters prove different things when they read a file: the cc adapter's
773/// input scan also establishes that a source embeds no timestamp macro, which
774/// a digest recorded by the rustc adapter never checked. Scoping the ledger
775/// keeps one adapter's shortcut from resting on a property another adapter
776/// never established.
777#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
778#[serde(rename_all = "snake_case")]
779pub enum FileDigestScope {
780    /// The digest describes the file's contents and nothing more.
781    Content,
782    /// The digest describes a cc compiler input that also passed the
783    /// timestamp-macro scan.
784    CcInput,
785}
786
787/// Recorded digests a session may consult instead of rehashing a file.
788///
789/// The agent's file-digest ledger answers through this everywhere a caller
790/// holds one; the no-op implementation stands in where reuse must not happen,
791/// such as under verification.
792pub trait FileDigestCache: Send + Sync {
793    /// Recorded digests for these identities, in request order.
794    fn find(&self, scope: FileDigestScope, files: &[FileIdentity]) -> Vec<Option<CacheDigest>>;
795    /// Record digests for files read under these identities.
796    fn record(&self, scope: FileDigestScope, entries: Vec<RecordedFileDigest>);
797}
798
799/// A [`FileDigestCache`] that remembers nothing and finds nothing.
800pub struct NoFileDigestCache;
801
802impl FileDigestCache for NoFileDigestCache {
803    fn find(&self, _scope: FileDigestScope, files: &[FileIdentity]) -> Vec<Option<CacheDigest>> {
804        vec![None; files.len()]
805    }
806
807    fn record(&self, _scope: FileDigestScope, _entries: Vec<RecordedFileDigest>) {}
808}
809
810/// A CAS blob whose contents this session hashed, and the file identity that
811/// says the hash still describes what is on disk.
812///
813/// Length alone would miss an overwrite that keeps the size; the modification
814/// time is what makes a rewrite visible without reading the bytes again.
815#[derive(Debug, Clone)]
816struct VerifiedBlob {
817    path: PathBuf,
818    len: u64,
819    modified: SystemTime,
820}
821
822impl VerifiedBlob {
823    /// Describe a file that was just verified, or nothing when the filesystem
824    /// will not report a modification time to compare against later.
825    fn describe(path: &Path) -> Option<Self> {
826        let metadata = std::fs::metadata(path).ok()?;
827        Some(Self {
828            path: path.to_path_buf(),
829            len: metadata.len(),
830            modified: metadata.modified().ok()?,
831        })
832    }
833
834    /// Whether the file still has the identity it had when it was verified.
835    fn is_unchanged(&self) -> bool {
836        let Ok(metadata) = std::fs::metadata(&self.path) else {
837            return false;
838        };
839        metadata.len() == self.len && metadata.modified().is_ok_and(|now| now == self.modified)
840    }
841}
842
843#[derive(Debug, Clone, Default)]
844struct TaskActionState {
845    manifest: String,
846    baseline_loaded: bool,
847    predictions: BTreeMap<CacheDigest, ActionPrediction>,
848    pending_predictions: BTreeMap<CacheDigest, ActionPrediction>,
849    remote_etag: Option<String>,
850}
851
852struct PrefetchedAction {
853    adapter: String,
854    result: RemoteActionResult,
855}
856
857struct RemoteDownloadReservation {
858    counter: Arc<AtomicU64>,
859    reserved: u64,
860    committed: bool,
861}
862
863impl RemoteDownloadReservation {
864    fn bytes(&self) -> u64 {
865        self.reserved
866    }
867
868    fn commit(mut self, bytes: u64) {
869        debug_assert!(bytes <= self.reserved);
870        self.counter
871            .fetch_sub(self.reserved.saturating_sub(bytes), Ordering::AcqRel);
872        self.committed = true;
873    }
874}
875
876impl Drop for RemoteDownloadReservation {
877    fn drop(&mut self) {
878        if !self.committed {
879            self.counter.fetch_sub(self.reserved, Ordering::AcqRel);
880        }
881    }
882}
883
884#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
885struct ExecutableIdentityKey {
886    executable: PathBuf,
887    environment: BTreeMap<String, Option<String>>,
888}
889
890impl CacheAgent {
891    /// Create an agent backed by the cache rooted at `cache_dir`.
892    pub fn new(cache_dir: impl Into<PathBuf>, version: impl Into<Arc<str>>) -> Self {
893        Self::build(cache_dir.into(), version.into(), None, 0)
894    }
895
896    /// Create an agent with local-first access to a remote action cache.
897    pub fn new_remote(
898        cache_dir: impl Into<PathBuf>,
899        version: impl Into<Arc<str>>,
900        remote: AgentRemoteCache,
901    ) -> Self {
902        Self::build(
903            cache_dir.into(),
904            version.into(),
905            Some(remote),
906            DEFAULT_MAX_REMOTE_DOWNLOAD_BYTES,
907        )
908    }
909
910    /// Create a remote agent with a cumulative download budget for the session.
911    pub fn new_remote_with_download_limit(
912        cache_dir: impl Into<PathBuf>,
913        version: impl Into<Arc<str>>,
914        remote: AgentRemoteCache,
915        max_remote_download_bytes: u64,
916    ) -> Self {
917        Self::build(
918            cache_dir.into(),
919            version.into(),
920            Some(remote),
921            max_remote_download_bytes,
922        )
923    }
924
925    fn build(
926        cache_dir: PathBuf,
927        version: Arc<str>,
928        remote: Option<AgentRemoteCache>,
929        remote_download_limit: u64,
930    ) -> Self {
931        let remote_mode = remote
932            .as_ref()
933            .map_or(RemoteCacheMode::ReadOnly, |remote| remote.mode);
934        let remote_staging_dir = remote.as_ref().map_or_else(
935            || cache_dir.join("remote"),
936            |remote| remote.staging_dir.clone(),
937        );
938        let remote = remote.map(|remote| Arc::new(remote.client));
939        let stats = Arc::new(AtomicAgentStats::default());
940        let remote_transfers = Arc::new(tokio::sync::Semaphore::new(MAX_REMOTE_TRANSFERS));
941        // A session that cannot publish never queues an upload, so it does not
942        // need somewhere to queue one.
943        let uploads = remote
944            .clone()
945            .filter(|_| remote_mode.writes())
946            .map(|client| {
947                UploadQueue::new(
948                    client,
949                    Arc::new(AgentUploadSink {
950                        stats: stats.clone(),
951                    }),
952                    remote_transfers.clone(),
953                )
954            });
955        Self {
956            cas: LocalCas::new(cache_dir.clone()),
957            actions: LocalActionCache::new(cache_dir.clone()),
958            verified_blobs: Arc::new(Mutex::new(BTreeMap::new())),
959            version,
960            write_locks: Arc::new(Mutex::new(BTreeMap::new())),
961            action_locks: Arc::new(Mutex::new(BTreeMap::new())),
962            stats,
963            observer: None,
964            executable_identities: Arc::new(Mutex::new(BTreeMap::new())),
965            manifest_dir: Arc::new(task_manifest_dir(&cache_dir)),
966            task_actions: Arc::new(Mutex::new(BTreeMap::new())),
967            next_task_run: Arc::new(AtomicU64::new(0)),
968            manifest_write_lock: Arc::new(Mutex::new(())),
969            remote,
970            remote_mode,
971            remote_staging_dir: Arc::new(remote_staging_dir),
972            remote_download_limit,
973            remote_download_bytes: Arc::new(AtomicU64::new(0)),
974            pending_remote_actions: Arc::new(Mutex::new(BTreeMap::new())),
975            remote_transfers,
976            prefetch_transfers: Arc::new(tokio::sync::Semaphore::new(MAX_PREFETCH_TRANSFERS)),
977            prefetch_tasks: Arc::new(Mutex::new(Vec::new())),
978            warnings: Arc::new(Mutex::new(BTreeSet::new())),
979            file_digests: Arc::new(Mutex::new(BTreeMap::new())),
980            uploads,
981        }
982    }
983
984    /// Report each accounted cache decision to `observer` as it happens.
985    #[must_use]
986    pub fn with_observer(mut self, observer: Arc<dyn AgentEventObserver>) -> Self {
987        self.observer = Some(observer);
988        self
989    }
990
991    fn emit(&self, event: impl FnOnce() -> AgentEvent) {
992        if let Some(observer) = &self.observer {
993            observer.event(event());
994        }
995    }
996
997    fn reserve_remote_download(&self, bytes: u64) -> Result<RemoteDownloadReservation> {
998        let mut current = self.remote_download_bytes.load(Ordering::Acquire);
999        loop {
1000            let next = current
1001                .checked_add(bytes)
1002                .ok_or_else(|| eyre::eyre!("remote cache download budget overflowed"))?;
1003            if next > self.remote_download_limit {
1004                bail!(
1005                    "remote cache download budget exceeded: {} bytes requested with {} of {} bytes already used",
1006                    bytes,
1007                    current,
1008                    self.remote_download_limit
1009                );
1010            }
1011            match self.remote_download_bytes.compare_exchange_weak(
1012                current,
1013                next,
1014                Ordering::AcqRel,
1015                Ordering::Acquire,
1016            ) {
1017                Ok(_) => {
1018                    return Ok(RemoteDownloadReservation {
1019                        counter: self.remote_download_bytes.clone(),
1020                        reserved: bytes,
1021                        committed: false,
1022                    });
1023                }
1024                Err(observed) => current = observed,
1025            }
1026        }
1027    }
1028
1029    fn reserve_remote_download_up_to(&self, requested: u64) -> Result<RemoteDownloadReservation> {
1030        let mut current = self.remote_download_bytes.load(Ordering::Acquire);
1031        loop {
1032            let reserved = requested.min(self.remote_download_limit.saturating_sub(current));
1033            let next = current + reserved;
1034            match self.remote_download_bytes.compare_exchange_weak(
1035                current,
1036                next,
1037                Ordering::AcqRel,
1038                Ordering::Acquire,
1039            ) {
1040                Ok(_) => {
1041                    return Ok(RemoteDownloadReservation {
1042                        counter: self.remote_download_bytes.clone(),
1043                        reserved,
1044                        committed: false,
1045                    });
1046                }
1047                Err(observed) => current = observed,
1048            }
1049        }
1050    }
1051
1052    /// Load the last committed action manifest for a task into this session.
1053    pub async fn begin_task(&self, task: &str) -> Result<String> {
1054        self.begin_task_with_remote_errors(task, false).await
1055    }
1056
1057    /// Load a task and finish its prefetch, surfacing remote lookup failures.
1058    pub async fn prefetch_task(&self, task: &str) -> Result<String> {
1059        let run = self.begin_task_with_remote_errors(task, true).await?;
1060        self.wait_for_prefetches().await;
1061        Ok(run)
1062    }
1063
1064    async fn begin_task_with_remote_errors(&self, task: &str, strict: bool) -> Result<String> {
1065        validate_task_identity(task)?;
1066        let (remote_manifest, mut remote_etag) = if self.remote_mode.reads() {
1067            match self.get_remote_task_manifest(task).await {
1068                Ok(Some((manifest, etag))) => (Some(manifest), Some(etag)),
1069                Ok(None) => (None, None),
1070                Err(error) => {
1071                    if strict {
1072                        return Err(error).wrap_err_with(|| {
1073                            format!("remote task action manifest lookup failed for {task}")
1074                        });
1075                    }
1076                    self.note_remote_failure();
1077                    warn!("remote task action manifest lookup failed for {task}: {error}");
1078                    (None, None)
1079                }
1080            }
1081        } else {
1082            (None, None)
1083        };
1084        let manifest = {
1085            let _write_guard = self.manifest_write_lock.lock().unwrap();
1086            let _file_guard = self.lock_task_manifest(task)?;
1087            let local_manifest = self.load_task_manifest(task)?;
1088            let manifest = match (remote_manifest, local_manifest) {
1089                (Some(remote), Some(local)) => {
1090                    let (manifest, merged) = merge_remote_task_manifest(task, remote, local);
1091                    if !merged {
1092                        remote_etag = None;
1093                    }
1094                    Some(manifest)
1095                }
1096                (Some(remote), None) => Some(remote),
1097                (None, local) => local,
1098            };
1099            if let Some(manifest) = &manifest {
1100                self.persist_task_manifest(manifest)?;
1101            }
1102            manifest
1103        };
1104        let state = if let Some(manifest) = manifest {
1105            TaskActionState {
1106                manifest: task.to_string(),
1107                baseline_loaded: true,
1108                predictions: manifest
1109                    .predictions
1110                    .into_iter()
1111                    .map(|prediction| (prediction.invocation.clone(), prediction))
1112                    .collect(),
1113                pending_predictions: BTreeMap::new(),
1114                remote_etag,
1115            }
1116        } else {
1117            TaskActionState {
1118                manifest: task.to_string(),
1119                baseline_loaded: true,
1120                remote_etag,
1121                ..TaskActionState::default()
1122            }
1123        };
1124        let sequence = self.next_task_run.fetch_add(1, Ordering::Relaxed);
1125        let run =
1126            CacheDigest::blake3(format!("{task}\0{}\0{sequence}", std::process::id()).as_bytes())
1127                .hash;
1128        // The largest baseline rather than a sum: beginning the same task again
1129        // in one session reloads the same manifest, and counting it twice would
1130        // overstate what there was to match.
1131        self.stats.predictions_loaded.fetch_max(
1132            state.predictions.len().try_into().unwrap_or(u64::MAX),
1133            Ordering::Relaxed,
1134        );
1135        let predictions = state.predictions.values().cloned().collect();
1136        self.task_actions.lock().unwrap().insert(run.clone(), state);
1137        self.spawn_prefetch_predictions(predictions);
1138        Ok(run)
1139    }
1140
1141    /// Cancel speculative downloads before the owning session exits.
1142    pub async fn cancel_prefetches(&self) {
1143        let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
1144        for task in &tasks {
1145            task.abort();
1146        }
1147        for task in tasks {
1148            if let Err(error) = task.await
1149                && !error.is_cancelled()
1150            {
1151                warn!("remote action prefetch task failed: {error}");
1152            }
1153        }
1154    }
1155
1156    /// Publish everything a build queued, before the session stops.
1157    ///
1158    /// Store requests return once an object is durable locally, so at this point
1159    /// the remote cache may still be behind the local one. Uploads run on the
1160    /// session's runtime and are abandoned if it goes away, so a session that
1161    /// wants them published has to wait here.
1162    pub async fn wait_for_uploads(&self) {
1163        let Some(uploads) = &self.uploads else {
1164            return;
1165        };
1166        let _timer = AtomicDurationTimer::start(&self.stats.upload_drain_duration_ns);
1167        uploads.drain().await;
1168    }
1169
1170    async fn wait_for_prefetches(&self) {
1171        let tasks = std::mem::take(&mut *self.prefetch_tasks.lock().unwrap());
1172        for task in tasks {
1173            if let Err(error) = task.await {
1174                warn!("remote action prefetch task failed: {error}");
1175            }
1176        }
1177    }
1178
1179    /// Atomically publish the completed actions collected by a task run.
1180    pub async fn commit_task(&self, run: &str) -> Result<()> {
1181        validate_task_identity(run)?;
1182        let state = self
1183            .task_actions
1184            .lock()
1185            .unwrap()
1186            .get(run)
1187            .cloned()
1188            .ok_or_else(|| eyre::eyre!("task action manifest baseline was not loaded"))?;
1189        if !state.baseline_loaded {
1190            bail!("task action manifest baseline was not loaded");
1191        }
1192        let task = state.manifest;
1193        validate_task_identity(&task)?;
1194        let (manifest, introduced) = {
1195            let _write_guard = self.manifest_write_lock.lock().unwrap();
1196            let _file_guard = self.lock_task_manifest(&task)?;
1197            let mut predictions = self
1198                .load_task_manifest(&task)?
1199                .map(|manifest| {
1200                    manifest
1201                        .predictions
1202                        .into_iter()
1203                        .map(|prediction| (prediction.invocation.clone(), prediction))
1204                        .collect::<BTreeMap<_, _>>()
1205                })
1206                .unwrap_or_default();
1207            // Which predictions this run adds, as opposed to inherits. Only a
1208            // new one may be withheld for a failed upload: retracting an
1209            // inherited one would un-advertise a result that is plausibly still
1210            // on the server from whichever session put it there.
1211            let introduced: BTreeSet<CacheDigest> = state
1212                .pending_predictions
1213                .keys()
1214                .filter(|invocation| !predictions.contains_key(*invocation))
1215                .cloned()
1216                .collect();
1217            // Only publish predictions recorded by this run. `predictions`
1218            // also contains the baseline loaded by `begin_task`; extending
1219            // with that snapshot would overwrite newer entries committed by
1220            // another agent process after this run began.
1221            predictions.extend(state.pending_predictions);
1222            let manifest = TaskActionManifest {
1223                version: TASK_ACTION_MANIFEST_VERSION,
1224                task: task.clone(),
1225                predictions: predictions.into_values().collect(),
1226            };
1227            validate_task_manifest(&manifest, &task)?;
1228            self.persist_task_manifest(&manifest)?;
1229            (manifest, introduced)
1230        };
1231        self.task_actions.lock().unwrap().remove(run);
1232        if self.remote_mode.writes() {
1233            // A manifest advertises the actions it predicts, so it must not
1234            // reach the remote cache before the results a reader would then go
1235            // looking for -- nor name a result that never got there at all.
1236            let mut manifest = manifest;
1237            if let Some(uploads) = &self.uploads {
1238                let actions: Vec<CacheDigest> = manifest
1239                    .predictions
1240                    .iter()
1241                    .map(|prediction| prediction.action.clone())
1242                    .collect();
1243                let unpublished = uploads.wait_for_actions(&actions).await;
1244                let withheld = manifest
1245                    .predictions
1246                    .iter()
1247                    .filter(|prediction| {
1248                        introduced.contains(&prediction.invocation)
1249                            && unpublished.contains(&prediction.action)
1250                    })
1251                    .count();
1252                if withheld > 0 {
1253                    // The local manifest keeps them: this checkout can still use
1254                    // what it built, and a later session can publish it.
1255                    warn!(
1256                        "{withheld} of {} predicted actions were not published, so the remote task action manifest omits them",
1257                        manifest.predictions.len()
1258                    );
1259                    manifest.predictions.retain(|prediction| {
1260                        !(introduced.contains(&prediction.invocation)
1261                            && unpublished.contains(&prediction.action))
1262                    });
1263                }
1264            }
1265            match self
1266                .put_remote_task_manifest(&task, manifest, state.remote_etag)
1267                .await
1268            {
1269                Ok(remote_manifest) => {
1270                    let _write_guard = self.manifest_write_lock.lock().unwrap();
1271                    let reconciliation = (|| {
1272                        let _file_guard = self.lock_task_manifest(&task)?;
1273                        let manifest = match self.load_task_manifest(&task)? {
1274                            Some(local) => {
1275                                merge_remote_task_manifest(&task, remote_manifest, local).0
1276                            }
1277                            None => remote_manifest,
1278                        };
1279                        self.persist_task_manifest(&manifest)
1280                    })();
1281                    if let Err(error) = reconciliation {
1282                        warn!(
1283                            "remote task action manifest reconciliation failed for {task}: {error}"
1284                        );
1285                    }
1286                }
1287                Err(error) => {
1288                    self.note_remote_failure();
1289                    warn!("remote task action manifest upload failed for {task}: {error}");
1290                }
1291            }
1292        }
1293        Ok(())
1294    }
1295
1296    fn task_manifest_path(&self, task: &str) -> PathBuf {
1297        self.manifest_dir.join(format!("{task}.json"))
1298    }
1299
1300    fn task_manifest_lock_path(&self, task: &str) -> PathBuf {
1301        self.manifest_dir.join("locks").join(format!("{task}.lock"))
1302    }
1303
1304    fn lock_task_manifest(&self, task: &str) -> Result<fslock::LockFile> {
1305        let path = self.task_manifest_lock_path(task);
1306        fs::create_dir_all(path.parent().expect("task manifest lock has a parent"))?;
1307        let mut lock = fslock::LockFile::open(&path)?;
1308        lock.lock()?;
1309        Ok(lock)
1310    }
1311
1312    fn load_task_manifest(&self, task: &str) -> Result<Option<TaskActionManifest>> {
1313        match fs::read(self.task_manifest_path(task)) {
1314            Ok(contents) => Ok(Some(self.parse_task_manifest(task, &contents, false)?)),
1315            Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
1316            Err(error) => Err(error.into()),
1317        }
1318    }
1319
1320    fn parse_task_manifest(
1321        &self,
1322        task: &str,
1323        contents: &[u8],
1324        require_canonical: bool,
1325    ) -> Result<TaskActionManifest> {
1326        let manifest: TaskActionManifest = serde_json::from_slice(contents)?;
1327        validate_task_manifest(&manifest, task)?;
1328        if require_canonical && canonical_json(&manifest)? != contents {
1329            bail!("task action manifest is not canonical JSON");
1330        }
1331        Ok(manifest)
1332    }
1333
1334    fn task_manifest_selector(task: &str) -> Result<(Vec<u8>, CacheDigest)> {
1335        TaskActionManifest::selector(task)
1336    }
1337
1338    fn persist_task_manifest(&self, manifest: &TaskActionManifest) -> Result<()> {
1339        let bytes = canonical_json(manifest)?;
1340        fs::create_dir_all(self.manifest_dir.as_path())?;
1341        let mut temporary = tempfile::NamedTempFile::new_in(self.manifest_dir.as_path())?;
1342        std::io::Write::write_all(temporary.as_file_mut(), &bytes)?;
1343        temporary.as_file_mut().sync_all()?;
1344        temporary
1345            .persist(self.task_manifest_path(&manifest.task))
1346            .map_err(|error| error.error)?;
1347        Ok(())
1348    }
1349
1350    async fn get_remote_task_manifest(
1351        &self,
1352        task: &str,
1353    ) -> Result<Option<(TaskActionManifest, String)>> {
1354        let Some(remote) = &self.remote else {
1355            return Ok(None);
1356        };
1357        let (_, selector) = Self::task_manifest_selector(task)?;
1358        let _permit = self.remote_transfers.acquire().await?;
1359        self.stats
1360            .remote_manifest_lookups
1361            .fetch_add(1, Ordering::Relaxed);
1362        let _timer = AtomicDurationTimer::start(&self.stats.remote_manifest_lookup_duration_ns);
1363        let Some(remote_manifest) = remote.get_action_manifest(&selector).await? else {
1364            return Ok(None);
1365        };
1366        let manifest = self.parse_task_manifest(task, &remote_manifest.bytes, true)?;
1367        Ok(Some((manifest, remote_manifest.etag)))
1368    }
1369
1370    async fn put_remote_task_manifest(
1371        &self,
1372        task: &str,
1373        mut manifest: TaskActionManifest,
1374        mut expected_etag: Option<String>,
1375    ) -> Result<TaskActionManifest> {
1376        let Some(remote) = &self.remote else {
1377            return Ok(manifest);
1378        };
1379        let (_, selector) = Self::task_manifest_selector(task)?;
1380        for _ in 0..4 {
1381            let bytes = canonical_json(&manifest)?;
1382            let outcome = {
1383                let _permit = self.remote_transfers.acquire().await?;
1384                remote
1385                    .put_action_manifest(&selector, &bytes, expected_etag.as_deref())
1386                    .await?
1387            };
1388            match outcome {
1389                ManifestPutOutcome::Stored => return Ok(manifest),
1390                ManifestPutOutcome::PreconditionFailed => {
1391                    let Some((remote_manifest, etag)) = self.get_remote_task_manifest(task).await?
1392                    else {
1393                        expected_etag = None;
1394                        continue;
1395                    };
1396                    manifest = merge_task_manifests(task, Some(remote_manifest), manifest)?;
1397                    expected_etag = Some(etag);
1398                }
1399            }
1400        }
1401        bail!("remote task action manifest changed too frequently")
1402    }
1403
1404    /// Return a snapshot of this session's cache activity.
1405    pub fn stats(&self) -> AgentStats {
1406        AgentStats {
1407            session_duration_ns: 0,
1408            lookups: self.stats.lookups.load(Ordering::Relaxed),
1409            unconsulted: self.stats.unconsulted.load(Ordering::Relaxed),
1410            hits: self.stats.hits.load(Ordering::Relaxed),
1411            stores: self.stats.stores.load(Ordering::Relaxed),
1412            stored_bytes: self.stats.stored_bytes.load(Ordering::Relaxed),
1413            verifications: self.stats.verifications.load(Ordering::Relaxed),
1414            divergences: self.stats.divergences.load(Ordering::Relaxed),
1415            downloaded_bytes: self.stats.downloaded_bytes.load(Ordering::Relaxed),
1416            uploaded_bytes: self.stats.uploaded_bytes.load(Ordering::Relaxed),
1417            background_uploads: self.stats.background_uploads.load(Ordering::Relaxed),
1418            background_upload_failures: self
1419                .stats
1420                .background_upload_failures
1421                .load(Ordering::Relaxed),
1422            remote_blob_pack_uploads: self.stats.remote_blob_pack_uploads.load(Ordering::Relaxed),
1423            remote_blob_pack_upload_blobs: self
1424                .stats
1425                .remote_blob_pack_upload_blobs
1426                .load(Ordering::Relaxed),
1427            upload_drain_duration_ns: self.stats.upload_drain_duration_ns.load(Ordering::Relaxed),
1428            prefetched_actions: self.stats.prefetched_actions.load(Ordering::Relaxed),
1429            predictions_loaded: self.stats.predictions_loaded.load(Ordering::Relaxed),
1430            bypasses: self.stats.bypasses.lock().unwrap().clone(),
1431            avoided_compiler_duration_ns: self
1432                .stats
1433                .avoided_compiler_duration_ns
1434                .load(Ordering::Relaxed),
1435            compiler: self.stats.compiler.lock().unwrap().clone(),
1436            slow_compilations: self.stats.slow_compilations.lock().unwrap().clone(),
1437            remote_failures: self.stats.remote_failures.load(Ordering::Relaxed),
1438            remote_manifest_lookups: self.stats.remote_manifest_lookups.load(Ordering::Relaxed),
1439            remote_manifest_lookup_duration_ns: self
1440                .stats
1441                .remote_manifest_lookup_duration_ns
1442                .load(Ordering::Relaxed),
1443            remote_action_lookups: self.stats.remote_action_lookups.load(Ordering::Relaxed),
1444            remote_action_lookup_duration_ns: self
1445                .stats
1446                .remote_action_lookup_duration_ns
1447                .load(Ordering::Relaxed),
1448            remote_blob_requests: self.stats.remote_blob_requests.load(Ordering::Relaxed),
1449            remote_blob_pack_requests: self.stats.remote_blob_pack_requests.load(Ordering::Relaxed),
1450            remote_blob_pack_blobs: self.stats.remote_blob_pack_blobs.load(Ordering::Relaxed),
1451            remote_blob_transfer_duration_ns: self
1452                .stats
1453                .remote_blob_transfer_duration_ns
1454                .load(Ordering::Relaxed),
1455            local_cas_write_duration_ns: self
1456                .stats
1457                .local_cas_write_duration_ns
1458                .load(Ordering::Relaxed),
1459            prefetch_runs: self.stats.prefetch_runs.load(Ordering::Relaxed),
1460            prefetch_duration_ns: self.stats.prefetch_duration_ns.load(Ordering::Relaxed),
1461            materialization_duration_ns: self
1462                .stats
1463                .materialization_duration_ns
1464                .load(Ordering::Relaxed),
1465            restored_output_files: self.stats.restored_output_files.load(Ordering::Relaxed),
1466            restored_output_bytes: self.stats.restored_output_bytes.load(Ordering::Relaxed),
1467            reflinked_output_files: self.stats.reflinked_output_files.load(Ordering::Relaxed),
1468            reflinked_output_bytes: self.stats.reflinked_output_bytes.load(Ordering::Relaxed),
1469            copied_output_files: self.stats.copied_output_files.load(Ordering::Relaxed),
1470            copied_output_bytes: self.stats.copied_output_bytes.load(Ordering::Relaxed),
1471            reused_output_files: self.stats.reused_output_files.load(Ordering::Relaxed),
1472            reused_output_bytes: self.stats.reused_output_bytes.load(Ordering::Relaxed),
1473        }
1474    }
1475
1476    /// Handle requests without a transport connection.
1477    ///
1478    /// Persistent wrappers use this entry point when Cargo invokes mbx outside
1479    /// an orchestrated session. It intentionally has the same response
1480    /// semantics as [`Self::handle_connection`], while leaving framing and the
1481    /// version handshake to callers that actually cross a process boundary.
1482    pub async fn handle_requests(
1483        &self,
1484        requests: impl IntoIterator<Item = AgentRequest>,
1485    ) -> Vec<AgentResponse> {
1486        let mut connection = ConnectionUploads::default();
1487        let mut responses = Vec::new();
1488        for request in requests {
1489            responses.push(self.respond_on(request, &mut connection).await);
1490        }
1491        responses
1492    }
1493
1494    fn write_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1495        Self::digest_lock(&self.write_locks, digest)
1496    }
1497
1498    fn action_lock(&self, digest: &CacheDigest) -> Arc<tokio::sync::Mutex<()>> {
1499        Self::digest_lock(&self.action_locks, digest)
1500    }
1501
1502    fn digest_lock(
1503        locks: &Mutex<BTreeMap<CacheDigest, Weak<tokio::sync::Mutex<()>>>>,
1504        digest: &CacheDigest,
1505    ) -> Arc<tokio::sync::Mutex<()>> {
1506        let mut locks = locks.lock().unwrap();
1507        locks.retain(|_, lock| lock.strong_count() > 0);
1508        if let Some(lock) = locks.get(digest).and_then(Weak::upgrade) {
1509            return lock;
1510        }
1511        let lock = Arc::new(tokio::sync::Mutex::new(()));
1512        locks.insert(digest.clone(), Arc::downgrade(&lock));
1513        lock
1514    }
1515
1516    fn spawn_prefetch_predictions(&self, predictions: Vec<ActionPrediction>) {
1517        if predictions.is_empty() || !self.remote_mode.reads() || self.remote.is_none() {
1518            return;
1519        }
1520        let agent = self.clone();
1521        let task = tokio::spawn(async move {
1522            agent.prefetch_predictions(predictions.iter()).await;
1523        });
1524        self.prefetch_tasks.lock().unwrap().push(task);
1525    }
1526
1527    async fn prefetch_predictions<'a>(
1528        &self,
1529        predictions: impl Iterator<Item = &'a ActionPrediction>,
1530    ) {
1531        if !self.remote_mode.reads() || self.remote.is_none() {
1532            return;
1533        }
1534        self.stats.prefetch_runs.fetch_add(1, Ordering::Relaxed);
1535        let _timer = AtomicDurationTimer::start(&self.stats.prefetch_duration_ns);
1536        let mut actions = BTreeMap::new();
1537        for prediction in predictions {
1538            actions
1539                .entry(prediction.action.clone())
1540                .or_insert_with(|| prediction.adapter.clone());
1541        }
1542        // One request per predicted action is the bulk of a prefetch's latency on
1543        // a large workspace, so ask for them together where the server allows it.
1544        match self.prefetch_action_batches(&actions).await {
1545            Ok(true) => return,
1546            Ok(false) => {}
1547            Err(error) => {
1548                self.note_remote_failure();
1549                warn!("remote action batch lookup failed: {error}");
1550            }
1551        }
1552        let mut actions = actions.into_iter();
1553        let mut tasks = tokio::task::JoinSet::new();
1554        for _ in 0..MAX_PREFETCH_TRANSFERS {
1555            let Some((action, adapter)) = actions.next() else {
1556                break;
1557            };
1558            let agent = self.clone();
1559            tasks.spawn(async move { agent.resolve_prefetch_action(action, adapter).await });
1560        }
1561        let mut resolved = Vec::new();
1562        while !tasks.is_empty() {
1563            let result = if resolved.is_empty() {
1564                tasks.join_next().await
1565            } else {
1566                match tokio::time::timeout(PREFETCH_ACTION_BATCH_DELAY, tasks.join_next()).await {
1567                    Ok(result) => result,
1568                    Err(_) => {
1569                        self.prefetch_resolved_actions(std::mem::take(&mut resolved))
1570                            .await;
1571                        continue;
1572                    }
1573                }
1574            };
1575            let Some(result) = result else {
1576                break;
1577            };
1578            match result {
1579                Ok(Ok(Some(action))) => resolved.push(action),
1580                Ok(Ok(None)) => {}
1581                Ok(Err(error)) => {
1582                    self.note_remote_failure();
1583                    warn!("remote action prefetch failed: {error}");
1584                }
1585                Err(error) => {
1586                    self.note_remote_failure();
1587                    warn!("remote action prefetch task failed: {error}");
1588                }
1589            }
1590            if let Some((action, adapter)) = actions.next() {
1591                let agent = self.clone();
1592                tasks.spawn(async move { agent.resolve_prefetch_action(action, adapter).await });
1593            }
1594            if resolved.len() == MAX_PREFETCH_ACTION_BATCH {
1595                self.prefetch_resolved_actions(std::mem::take(&mut resolved))
1596                    .await;
1597            }
1598        }
1599        if !resolved.is_empty() {
1600            self.prefetch_resolved_actions(resolved).await;
1601        }
1602    }
1603
1604    /// Resolve predicted actions in batched lookups, staging what comes back.
1605    ///
1606    /// Returns whether the batch extension answered. A server without it, or one
1607    /// that stops answering part way through, leaves the per-action path to
1608    /// resolve whatever is left -- the results already staged here are memoized,
1609    /// so nothing is looked up twice.
1610    async fn prefetch_action_batches(
1611        &self,
1612        actions: &BTreeMap<CacheDigest, String>,
1613    ) -> Result<bool> {
1614        let Some(remote) = self.remote.as_deref() else {
1615            return Ok(false);
1616        };
1617        let Some(limit) = remote.action_batch_limit().await? else {
1618            return Ok(false);
1619        };
1620        let wanted: Vec<CacheDigest> = actions
1621            .keys()
1622            .filter(|action| !self.action_is_staged(action))
1623            .cloned()
1624            .collect();
1625        if wanted.is_empty() {
1626            return Ok(true);
1627        }
1628        let mut chunks: Vec<Vec<CacheDigest>> = Vec::new();
1629        for chunk in wanted.chunks(limit) {
1630            chunks.push(chunk.to_vec());
1631        }
1632        let mut lookups = stream::iter(chunks)
1633            .map(|chunk| async move {
1634                let _prefetch_permit = self.prefetch_transfers.acquire().await?;
1635                let _permit = self.remote_transfers.acquire().await?;
1636                self.stats
1637                    .remote_action_lookups
1638                    .fetch_add(1, Ordering::Relaxed);
1639                let _timer =
1640                    AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
1641                remote.get_action_results(&chunk).await
1642            })
1643            .buffer_unordered(MAX_PREFETCH_BATCH_LOOKUPS);
1644        let mut answered = true;
1645        while let Some(lookup) = lookups.next().await {
1646            let results = match lookup {
1647                Ok(Some(results)) => results,
1648                // Advertised but unavailable. What is left falls back.
1649                Ok(None) => {
1650                    answered = false;
1651                    continue;
1652                }
1653                Err(error) => {
1654                    self.note_remote_failure();
1655                    warn!("remote action batch lookup failed: {error}");
1656                    answered = false;
1657                    continue;
1658                }
1659            };
1660            let mut resolved = Vec::with_capacity(results.len());
1661            for result in results {
1662                let Some(adapter) = actions.get(&result.action).cloned() else {
1663                    continue;
1664                };
1665                let lock = self.action_lock(&result.action);
1666                let _guard = lock.lock().await;
1667                if self.actions.find(&result.action)?.is_some() {
1668                    continue;
1669                }
1670                self.pending_remote_actions
1671                    .lock()
1672                    .unwrap()
1673                    .insert(result.action.clone(), result.clone());
1674                resolved.push(PrefetchedAction { adapter, result });
1675            }
1676            while !resolved.is_empty() {
1677                let wave = resolved
1678                    .drain(..resolved.len().min(MAX_PREFETCH_ACTION_BATCH))
1679                    .collect();
1680                self.prefetch_resolved_actions(wave).await;
1681            }
1682        }
1683        Ok(answered)
1684    }
1685
1686    /// Whether an action's result is already local or already looked up.
1687    fn action_is_staged(&self, action: &CacheDigest) -> bool {
1688        if self
1689            .pending_remote_actions
1690            .lock()
1691            .unwrap()
1692            .contains_key(action)
1693        {
1694            return true;
1695        }
1696        self.actions.find(action).is_ok_and(|found| found.is_some())
1697    }
1698
1699    #[cfg(test)]
1700    async fn prefetch_action(&self, action: CacheDigest, adapter: String) -> Result<()> {
1701        if let Some(action) = self.resolve_prefetch_action(action, adapter).await? {
1702            self.prefetch_resolved_actions(vec![action]).await;
1703        }
1704        Ok(())
1705    }
1706
1707    async fn resolve_prefetch_action(
1708        &self,
1709        action: CacheDigest,
1710        adapter: String,
1711    ) -> Result<Option<PrefetchedAction>> {
1712        let remote = self
1713            .remote
1714            .as_ref()
1715            .ok_or_else(|| eyre::eyre!("remote cache is not configured"))?;
1716        let result = {
1717            let lock = self.action_lock(&action);
1718            let _guard = lock.lock().await;
1719            if self.actions.find(&action)?.is_some() {
1720                return Ok(None);
1721            }
1722            if let Some(result) = self
1723                .pending_remote_actions
1724                .lock()
1725                .unwrap()
1726                .get(&action)
1727                .cloned()
1728            {
1729                result
1730            } else {
1731                let _prefetch_permit = self.prefetch_transfers.acquire().await?;
1732                let result = {
1733                    let _permit = self.remote_transfers.acquire().await?;
1734                    self.get_remote_action_result(remote, &action).await?
1735                };
1736                let Some(result) = result else {
1737                    return Ok(None);
1738                };
1739                self.pending_remote_actions
1740                    .lock()
1741                    .unwrap()
1742                    .insert(action.clone(), result.clone());
1743                result
1744            }
1745        };
1746        Ok(Some(PrefetchedAction { adapter, result }))
1747    }
1748
1749    fn prefetch_resolved_actions(&self, actions: Vec<PrefetchedAction>) -> BoxFuture<'_, ()> {
1750        self.prefetch_resolved_actions_inner(actions).boxed()
1751    }
1752
1753    async fn prefetch_resolved_actions_inner(&self, actions: Vec<PrefetchedAction>) {
1754        let Some(remote) = self.remote.as_deref() else {
1755            return;
1756        };
1757        if actions.is_empty() {
1758            return;
1759        }
1760
1761        let mut top_level = BTreeMap::new();
1762        for action in &actions {
1763            for digest in [
1764                Some(&action.result.action),
1765                action.result.metadata.as_ref(),
1766                action.result.output_root.as_ref(),
1767            ]
1768            .into_iter()
1769            .flatten()
1770            {
1771                top_level.insert(digest.clone(), ());
1772            }
1773        }
1774        let mut verified = self
1775            .fetch_remote_blobs(
1776                remote,
1777                top_level.into_keys().collect(),
1778                Some(&self.prefetch_transfers),
1779            )
1780            .await;
1781
1782        let mut next = BTreeMap::new();
1783        let mut pending_directories = BTreeMap::new();
1784        let mut parsed_directories = BTreeMap::new();
1785        let mut rustc_metadata = BTreeMap::new();
1786        for action in &actions {
1787            if action.adapter == "rustc"
1788                && let Some(metadata_digest) = &action.result.metadata
1789            {
1790                match verified
1791                    .get(metadata_digest)
1792                    .ok_or_else(|| eyre::eyre!("remote rustc action metadata is missing"))
1793                    .and_then(|path| Self::parse_rustc_metadata(path))
1794                {
1795                    Ok(metadata) => {
1796                        queue_prefetch_digest(&verified, &mut next, metadata.stdout.clone());
1797                        queue_prefetch_digest(&verified, &mut next, metadata.stderr.clone());
1798                        rustc_metadata.insert(metadata_digest.clone(), metadata);
1799                    }
1800                    Err(error) => warn!(
1801                        "remote rustc action metadata prefetch failed for {}: {error}",
1802                        action.result.action.hash
1803                    ),
1804                }
1805            }
1806            if let Some(output_root) = &action.result.output_root {
1807                pending_directories.insert(output_root.clone(), ());
1808            }
1809        }
1810
1811        let mut seen_directories = BTreeMap::new();
1812        loop {
1813            let mut following = BTreeMap::new();
1814            let mut directory_limit_exceeded = false;
1815            for digest in pending_directories.into_keys() {
1816                following.remove(&digest);
1817                if seen_directories.insert(digest.clone(), ()).is_some() {
1818                    continue;
1819                }
1820                if seen_directories.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
1821                    warn!("remote action output tree is too large to prefetch");
1822                    following.clear();
1823                    break;
1824                }
1825                match verified
1826                    .get(&digest)
1827                    .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))
1828                    .and_then(|path| Self::parse_cache_directory(path))
1829                {
1830                    Ok(directory) => {
1831                        for file in &directory.files {
1832                            queue_prefetch_digest(&verified, &mut next, file.digest.clone());
1833                            if next.len() >= MAX_PREFETCH_OBJECTS_PER_WAVE {
1834                                self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1835                                    .await;
1836                            }
1837                        }
1838                        for child in &directory.directories {
1839                            if !queue_prefetch_directory(
1840                                &seen_directories,
1841                                &mut following,
1842                                child.digest.clone(),
1843                                MAX_PREFETCH_DIRECTORY_OBJECTS,
1844                            ) {
1845                                warn!("remote action output tree is too large to prefetch");
1846                                directory_limit_exceeded = true;
1847                                break;
1848                            }
1849                            queue_prefetch_digest(&verified, &mut next, child.digest.clone());
1850                            if next.len() >= MAX_PREFETCH_OBJECTS_PER_WAVE {
1851                                self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1852                                    .await;
1853                            }
1854                        }
1855                        parsed_directories.insert(digest, directory);
1856                    }
1857                    Err(error) => warn!(
1858                        "remote action output directory prefetch failed for {}: {error}",
1859                        digest.hash
1860                    ),
1861                }
1862                if directory_limit_exceeded {
1863                    following.clear();
1864                    break;
1865                }
1866            }
1867            self.flush_prefetch_digest_batch(remote, &mut verified, &mut next)
1868                .await;
1869            if following.is_empty() {
1870                break;
1871            }
1872            pending_directories = following;
1873        }
1874
1875        for action in actions {
1876            match Self::validate_prefetched_action(
1877                &action,
1878                &verified,
1879                &rustc_metadata,
1880                &parsed_directories,
1881            ) {
1882                Ok(()) => {
1883                    if let Err(error) = self.actions.store(&action.result) {
1884                        warn!(
1885                            "remote action prefetch could not publish {}: {error}",
1886                            action.result.action.hash
1887                        );
1888                        continue;
1889                    }
1890                    self.pending_remote_actions
1891                        .lock()
1892                        .unwrap()
1893                        .remove(&action.result.action);
1894                    self.stats
1895                        .prefetched_actions
1896                        .fetch_add(1, Ordering::Relaxed);
1897                }
1898                Err(error) => warn!(
1899                    "remote action prefetch was incomplete for {}: {error}",
1900                    action.result.action.hash
1901                ),
1902            }
1903        }
1904    }
1905
1906    async fn flush_prefetch_digest_batch(
1907        &self,
1908        remote: &RemoteCacheClient,
1909        verified: &mut BTreeMap<CacheDigest, PathBuf>,
1910        pending: &mut BTreeMap<CacheDigest, ()>,
1911    ) {
1912        if pending.is_empty() {
1913            return;
1914        }
1915        let digests = std::mem::take(pending).into_keys().collect();
1916        verified.extend(
1917            self.fetch_remote_blobs(remote, digests, Some(&self.prefetch_transfers))
1918                .await,
1919        );
1920    }
1921
1922    async fn fetch_remote_blobs(
1923        &self,
1924        remote: &RemoteCacheClient,
1925        digests: Vec<CacheDigest>,
1926        prefetch_limit: Option<&tokio::sync::Semaphore>,
1927    ) -> BTreeMap<CacheDigest, PathBuf> {
1928        let mut verified = BTreeMap::new();
1929        let mut missing = BTreeMap::new();
1930        for digest in digests {
1931            match self.find_verified_blob(&digest) {
1932                Ok(Some(path)) => {
1933                    verified.insert(digest, path);
1934                }
1935                Ok(None) => {
1936                    missing.insert(digest, ());
1937                }
1938                Err(error) => warn!(
1939                    "local cache blob lookup failed for {}: {error}",
1940                    digest.hash
1941                ),
1942            }
1943        }
1944        if missing.is_empty() {
1945            return verified;
1946        }
1947
1948        let mut pack_candidates = missing.clone();
1949        while !pack_candidates.is_empty() {
1950            let candidates = match blob_pack_chunk(
1951                &pack_candidates.keys().cloned().collect::<Vec<_>>(),
1952                BlobPackLimits {
1953                    max_items: MAX_STAGED_BLOB_PACK_ITEMS,
1954                    max_bytes: MAX_STAGED_BLOB_PACK_BYTES,
1955                },
1956            ) {
1957                Ok(candidates) if !candidates.is_empty() => candidates,
1958                Ok(_) => break,
1959                Err(error) => {
1960                    warn!("remote cache blob pack skipped: {error}");
1961                    break;
1962                }
1963            };
1964            // A pack and an individual fetch share these per-digest locks. Hold
1965            // them through ingestion so overlapping prefetch and foreground
1966            // requests cannot download or charge the same object twice.
1967            let mut pack_guards = BTreeMap::new();
1968            for digest in candidates {
1969                let guard = self.write_lock(&digest).lock_owned().await;
1970                match self.find_verified_blob(&digest) {
1971                    Ok(Some(path)) => {
1972                        pack_candidates.remove(&digest);
1973                        missing.remove(&digest);
1974                        verified.insert(digest, path);
1975                    }
1976                    Ok(None) => {
1977                        pack_guards.insert(digest, guard);
1978                    }
1979                    Err(error) => {
1980                        warn!(
1981                            "local cache blob lookup failed for {}: {error}",
1982                            digest.hash
1983                        );
1984                        pack_guards.insert(digest, guard);
1985                    }
1986                }
1987            }
1988            let requested = pack_guards.keys().cloned().collect::<Vec<_>>();
1989            if requested.is_empty() {
1990                continue;
1991            }
1992            let requested_bytes = requested
1993                .iter()
1994                .fold(0_u64, |total, digest| total.saturating_add(digest.size));
1995            let pack_reservation = match self
1996                .reserve_remote_download_up_to(requested_bytes.min(MAX_STAGED_BLOB_PACK_BYTES))
1997            {
1998                Ok(reservation) if reservation.bytes() > 0 => reservation,
1999                Ok(_) => break,
2000                Err(error) => {
2001                    warn!("remote cache blob pack skipped: {error}");
2002                    break;
2003                }
2004            };
2005            let (pack, transfer_duration_ns) = {
2006                let _prefetch_permit = match prefetch_limit {
2007                    Some(limit) => match limit.acquire().await {
2008                        Ok(permit) => Some(permit),
2009                        Err(error) => {
2010                            warn!(
2011                                "remote cache blob pack could not acquire prefetch limit: {error}"
2012                            );
2013                            break;
2014                        }
2015                    },
2016                    None => None,
2017                };
2018                let _transfer_permit = match self.remote_transfers.acquire().await {
2019                    Ok(permit) => permit,
2020                    Err(error) => {
2021                        warn!("remote cache blob pack could not acquire transfer limit: {error}");
2022                        break;
2023                    }
2024                };
2025                let transfer_started = Instant::now();
2026                let pack = remote
2027                    .get_blob_pack_with_limit(
2028                        &requested,
2029                        self.remote_staging_dir.as_path(),
2030                        pack_reservation.bytes(),
2031                    )
2032                    .await;
2033                (pack, duration_ns(transfer_started))
2034            };
2035            let pack = match pack {
2036                Ok(Some(pack)) => pack,
2037                Ok(None) => break,
2038                Err(error) => {
2039                    atomic_saturating_add(
2040                        &self.stats.remote_blob_transfer_duration_ns,
2041                        transfer_duration_ns,
2042                    );
2043                    warn!(
2044                        "remote cache blob pack failed; falling back to individual blobs: {error}"
2045                    );
2046                    break;
2047                }
2048            };
2049            pack_reservation.commit(pack.payload_bytes);
2050            atomic_saturating_add(
2051                &self.stats.remote_blob_transfer_duration_ns,
2052                transfer_duration_ns,
2053            );
2054            atomic_saturating_add(&self.stats.remote_blob_pack_requests, pack.requests);
2055            atomic_saturating_add(&self.stats.remote_blob_pack_blobs, pack.blob_count);
2056            atomic_saturating_add(&self.stats.downloaded_bytes, pack.payload_bytes);
2057            if pack.requested.is_empty() {
2058                // The server's negotiated cap can be smaller than the local
2059                // staging cap used to select this locked slice. Fall back for
2060                // this slice, but keep packing later candidates that may fit.
2061                for digest in &requested {
2062                    pack_candidates.remove(digest);
2063                }
2064                continue;
2065            }
2066            for digest in &pack.requested {
2067                pack_candidates.remove(digest);
2068            }
2069            let mut ingests = stream::iter(pack.blobs.into_iter().map(|(digest, source)| {
2070                let digest_for_result = digest.clone();
2071                let guard = pack_guards
2072                    .remove(&digest)
2073                    .expect("requested packed blob has a write lock");
2074                async move {
2075                    (
2076                        digest_for_result,
2077                        self.ingest_packed_blob(digest, source, guard).await,
2078                    )
2079                }
2080            }))
2081            .buffer_unordered(MAX_PREFETCH_TRANSFERS);
2082            while let Some((digest, result)) = ingests.next().await {
2083                match result {
2084                    Ok(path) => {
2085                        missing.remove(&digest);
2086                        verified.insert(digest, path);
2087                    }
2088                    Err(error) => warn!(
2089                        "remote cache packed blob ingest failed for {}: {error}",
2090                        digest.hash
2091                    ),
2092                }
2093            }
2094        }
2095
2096        let mut transfers = stream::iter(missing.into_keys().map(|digest| {
2097            let digest_for_result = digest.clone();
2098            async move {
2099                (
2100                    digest_for_result,
2101                    self.fetch_remote_blob_with_limit(remote, &digest, prefetch_limit)
2102                        .await,
2103                )
2104            }
2105        }))
2106        .buffer_unordered(MAX_PREFETCH_TRANSFERS);
2107        while let Some((digest, result)) = transfers.next().await {
2108            match result {
2109                Ok(path) => {
2110                    verified.insert(digest, path);
2111                }
2112                Err(error) => warn!(
2113                    "remote cache blob prefetch failed for {}: {error}",
2114                    digest.hash
2115                ),
2116            }
2117        }
2118        verified
2119    }
2120
2121    async fn ingest_packed_blob(
2122        &self,
2123        digest: CacheDigest,
2124        source: PathBuf,
2125        _guard: tokio::sync::OwnedMutexGuard<()>,
2126    ) -> Result<PathBuf> {
2127        let digest_size = digest.size;
2128        let agent = self.clone();
2129        let (path, stored, cas_duration_ns) = tokio::task::spawn_blocking(move || {
2130            if let Some(path) = agent.find_verified_blob(&digest)? {
2131                return Ok::<_, eyre::Report>((path, false, 0));
2132            }
2133            let cas_started = Instant::now();
2134            let path = agent.cas.store_verified_file(&digest, &source)?;
2135            let cas_duration_ns = duration_ns(cas_started);
2136            agent.remember_verified_blob(&digest, &path);
2137            Ok((path, true, cas_duration_ns))
2138        })
2139        .await??;
2140        atomic_saturating_add(&self.stats.local_cas_write_duration_ns, cas_duration_ns);
2141        if stored {
2142            self.stats.stores.fetch_add(1, Ordering::Relaxed);
2143            atomic_saturating_add(&self.stats.stored_bytes, digest_size);
2144        }
2145        Ok(path)
2146    }
2147
2148    fn parse_rustc_metadata(path: &Path) -> Result<RustcMetadata> {
2149        let bytes = fs::read(path)?;
2150        let metadata: RustcMetadata = serde_json::from_slice(&bytes)?;
2151        if metadata.version != 1 || metadata.kind != "rustc" || canonical_json(&metadata)? != bytes
2152        {
2153            bail!("remote rustc action metadata is invalid");
2154        }
2155        Ok(metadata)
2156    }
2157
2158    fn parse_cache_directory(path: &Path) -> Result<CacheDirectory> {
2159        let bytes = fs::read(path)?;
2160        let directory: CacheDirectory = serde_json::from_slice(&bytes)?;
2161        if directory.version != 1 || canonical_json(&directory)? != bytes {
2162            bail!("remote action output directory is invalid");
2163        }
2164        Ok(directory)
2165    }
2166
2167    #[cfg(test)]
2168    fn load_cache_directory(&self, digest: &CacheDigest) -> Result<CacheDirectory> {
2169        let path = self
2170            .find_verified_blob(digest)?
2171            .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))?;
2172        Self::parse_cache_directory(&path)
2173    }
2174
2175    fn validate_prefetched_action(
2176        action: &PrefetchedAction,
2177        verified: &BTreeMap<CacheDigest, PathBuf>,
2178        rustc_metadata: &BTreeMap<CacheDigest, RustcMetadata>,
2179        directories: &BTreeMap<CacheDigest, CacheDirectory>,
2180    ) -> Result<()> {
2181        if !verified.contains_key(&action.result.action) {
2182            bail!("remote action descriptor is missing");
2183        }
2184        if let Some(metadata) = &action.result.metadata {
2185            if action.adapter == "rustc" {
2186                let metadata = rustc_metadata
2187                    .get(metadata)
2188                    .ok_or_else(|| eyre::eyre!("remote rustc action metadata is missing"))?;
2189                for digest in [&metadata.stdout, &metadata.stderr] {
2190                    if !verified.contains_key(digest) {
2191                        bail!("remote rustc action diagnostic blob is missing");
2192                    }
2193                }
2194            } else if !verified.contains_key(metadata) {
2195                bail!("remote action metadata is missing");
2196            }
2197        }
2198        let mut pending = action
2199            .result
2200            .output_root
2201            .iter()
2202            .cloned()
2203            .collect::<Vec<_>>();
2204        let mut seen = BTreeMap::new();
2205        while let Some(digest) = pending.pop() {
2206            if seen.insert(digest.clone(), ()).is_some() {
2207                continue;
2208            }
2209            if seen.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
2210                bail!("remote action output tree is too large");
2211            }
2212            let directory = directories
2213                .get(&digest)
2214                .ok_or_else(|| eyre::eyre!("remote action output directory is missing"))?;
2215            for file in &directory.files {
2216                if !verified.contains_key(&file.digest) {
2217                    bail!("remote action output file is missing");
2218                }
2219            }
2220            pending.extend(
2221                directory
2222                    .directories
2223                    .iter()
2224                    .map(|directory| directory.digest.clone()),
2225            );
2226        }
2227        Ok(())
2228    }
2229
2230    #[cfg(test)]
2231    async fn prefetch_output_tree(
2232        &self,
2233        remote: &RemoteCacheClient,
2234        output_root: &CacheDigest,
2235    ) -> Result<()> {
2236        let mut pending = vec![output_root.clone()];
2237        let mut seen = BTreeMap::new();
2238        while let Some(digest) = pending.pop() {
2239            if seen.insert(digest.clone(), ()).is_some() {
2240                continue;
2241            }
2242            if seen.len() > MAX_PREFETCH_DIRECTORY_OBJECTS {
2243                bail!("remote action output tree is too large");
2244            }
2245            self.fetch_remote_blob_with_limit(remote, &digest, Some(&self.prefetch_transfers))
2246                .await?;
2247            let directory = self.load_cache_directory(&digest)?;
2248            let mut transfers = stream::iter(directory.files.into_iter().map(|file| async move {
2249                self.fetch_remote_blob_with_limit(
2250                    remote,
2251                    &file.digest,
2252                    Some(&self.prefetch_transfers),
2253                )
2254                .await
2255                .map(|_| ())
2256            }))
2257            .buffer_unordered(MAX_PREFETCH_TRANSFERS);
2258            while let Some(result) = transfers.next().await {
2259                result?;
2260            }
2261            pending.extend(
2262                directory
2263                    .directories
2264                    .into_iter()
2265                    .map(|directory| directory.digest),
2266            );
2267        }
2268        Ok(())
2269    }
2270
2271    async fn fetch_remote_blob(
2272        &self,
2273        remote: &RemoteCacheClient,
2274        digest: &CacheDigest,
2275    ) -> Result<PathBuf> {
2276        self.fetch_remote_blob_with_limit(remote, digest, None)
2277            .await
2278    }
2279
2280    async fn fetch_remote_blob_with_limit(
2281        &self,
2282        remote: &RemoteCacheClient,
2283        digest: &CacheDigest,
2284        prefetch_limit: Option<&tokio::sync::Semaphore>,
2285    ) -> Result<PathBuf> {
2286        let lock = self.write_lock(digest);
2287        let _guard = lock.lock().await;
2288        if let Some(path) = self.find_verified_blob(digest)? {
2289            return Ok(path);
2290        }
2291        let _prefetch_permit = match prefetch_limit {
2292            Some(limit) => Some(limit.acquire().await?),
2293            None => None,
2294        };
2295        let _permit = self.remote_transfers.acquire().await?;
2296        let reservation = self.reserve_remote_download(digest.size)?;
2297        self.stats
2298            .remote_blob_requests
2299            .fetch_add(1, Ordering::Relaxed);
2300        let transfer_timer =
2301            AtomicDurationTimer::start(&self.stats.remote_blob_transfer_duration_ns);
2302        let temporary = remote
2303            .get_blob_file(digest, self.remote_staging_dir.as_path())
2304            .await?;
2305        drop(transfer_timer);
2306        let _cas_timer = AtomicDurationTimer::start(&self.stats.local_cas_write_duration_ns);
2307        let path = self.cas.store_verified_file(digest, temporary.path())?;
2308        reservation.commit(digest.size);
2309        self.remember_verified_blob(digest, &path);
2310        self.stats.stores.fetch_add(1, Ordering::Relaxed);
2311        self.stats
2312            .stored_bytes
2313            .fetch_add(digest.size, Ordering::Relaxed);
2314        self.stats
2315            .downloaded_bytes
2316            .fetch_add(digest.size, Ordering::Relaxed);
2317        Ok(path)
2318    }
2319
2320    /// Answer one request outside any connection.
2321    ///
2322    /// An upload queued this way has no sibling requests to order against, so
2323    /// the queue falls back to treating every blob it holds as a prerequisite.
2324    #[cfg(test)]
2325    async fn respond(&self, request: AgentRequest) -> AgentResponse {
2326        self.respond_on(request, &mut ConnectionUploads::default())
2327            .await
2328    }
2329
2330    async fn respond_on(
2331        &self,
2332        request: AgentRequest,
2333        connection: &mut ConnectionUploads,
2334    ) -> AgentResponse {
2335        let result = match request {
2336            AgentRequest::BeginTask { task } => self
2337                .begin_task(&task)
2338                .await
2339                .map(|run| AgentResponse::TaskBegun { run }),
2340            AgentRequest::CommitTask { run } => self
2341                .commit_task(&run)
2342                .await
2343                .map(|()| AgentResponse::TaskCommitted),
2344            AgentRequest::FindBlob { digest } => self.find_blob(&digest).await,
2345            AgentRequest::FindBlobs { digests } => self.find_blobs(digests).await,
2346            AgentRequest::StoreBlob { digest, source } => {
2347                self.store_blob(&digest, &source, connection).await
2348            }
2349            AgentRequest::FindActionResult { action } => {
2350                self.stats.lookups.fetch_add(1, Ordering::Relaxed);
2351                self.find_action_result(&action).await
2352            }
2353            AgentRequest::RecordActionHit {
2354                action,
2355                restore,
2356                crate_name,
2357            } => self.record_action_hit(&action, restore, crate_name),
2358            AgentRequest::RecordBypass { kind } => {
2359                *self
2360                    .stats
2361                    .bypasses
2362                    .lock()
2363                    .unwrap()
2364                    .entry(kind.clone())
2365                    .or_insert(0) += 1;
2366                self.emit(|| AgentEvent::Bypass { kind });
2367                Ok(AgentResponse::BypassRecorded)
2368            }
2369            AgentRequest::RecordUnconsulted => {
2370                self.stats.unconsulted.fetch_add(1, Ordering::Relaxed);
2371                self.emit(|| AgentEvent::Unconsulted);
2372                Ok(AgentResponse::UnconsultedRecorded)
2373            }
2374            AgentRequest::RecordWarning { message } => self.record_warning(message),
2375            AgentRequest::FindFileDigests { scope, files } => self.find_file_digests(scope, files),
2376            AgentRequest::RecordFileDigests { scope, entries } => {
2377                self.record_file_digests(scope, entries)
2378            }
2379            AgentRequest::RecordCompilerInvocation {
2380                outcome,
2381                crate_name,
2382                duration_ns,
2383            } => self.record_compiler_invocation(&outcome, crate_name.as_deref(), duration_ns),
2384            AgentRequest::RecordActionVerification { matched, restore } => {
2385                self.record_materialization(restore);
2386                self.stats.verifications.fetch_add(1, Ordering::Relaxed);
2387                if !matched {
2388                    self.stats.divergences.fetch_add(1, Ordering::Relaxed);
2389                }
2390                self.emit(|| AgentEvent::Verification { matched, restore });
2391                Ok(AgentResponse::ActionVerificationRecorded)
2392            }
2393            AgentRequest::StoreActionResult { result } => {
2394                self.store_action_result(&result, connection).await
2395            }
2396            AgentRequest::FindActionPrediction { task, invocation } => {
2397                self.find_action_prediction(&task, &invocation)
2398            }
2399            AgentRequest::RecordActionPrediction { task, prediction } => {
2400                self.record_action_prediction(&task, prediction)
2401            }
2402            AgentRequest::FindExecutableIdentity {
2403                executable,
2404                environment,
2405            } => self.find_executable_identity(executable, environment),
2406            AgentRequest::StoreExecutableIdentity {
2407                executable,
2408                environment,
2409                stdout,
2410            } => self.store_executable_identity(executable, environment, stdout),
2411            AgentRequest::Hello { .. } => {
2412                Err(eyre::eyre!("hello is only valid as the first request"))
2413            }
2414        };
2415        result.unwrap_or_else(|error| AgentResponse::Error {
2416            message: error.to_string(),
2417        })
2418    }
2419
2420    async fn find_blob(&self, digest: &CacheDigest) -> Result<AgentResponse> {
2421        if let Some(path) = self.find_verified_blob(digest)? {
2422            return Ok(AgentResponse::Blob { path: Some(path) });
2423        }
2424        if !self.remote_mode.reads() {
2425            return Ok(AgentResponse::Blob { path: None });
2426        }
2427        let Some(remote) = &self.remote else {
2428            return Ok(AgentResponse::Blob { path: None });
2429        };
2430        match self.fetch_remote_blob(remote, digest).await {
2431            Ok(path) => Ok(AgentResponse::Blob { path: Some(path) }),
2432            Err(error) => {
2433                warn!(
2434                    "remote cache blob lookup failed for {}: {error}",
2435                    digest.hash
2436                );
2437                Ok(AgentResponse::Blob { path: None })
2438            }
2439        }
2440    }
2441
2442    async fn find_blobs(&self, digests: Vec<CacheDigest>) -> Result<AgentResponse> {
2443        let mut paths = BTreeMap::new();
2444        let mut missing = Vec::new();
2445        for digest in &digests {
2446            match self.find_verified_blob(digest)? {
2447                Some(path) => {
2448                    paths.insert(digest.clone(), path);
2449                }
2450                None => {
2451                    missing.push(digest.clone());
2452                }
2453            }
2454        }
2455
2456        if !missing.is_empty()
2457            && self.remote_mode.reads()
2458            && let Some(remote) = &self.remote
2459        {
2460            paths.extend(self.fetch_remote_blobs(remote, missing, None).await);
2461        }
2462
2463        Ok(AgentResponse::Blobs {
2464            paths: digests
2465                .into_iter()
2466                .map(|digest| paths.get(&digest).cloned())
2467                .collect(),
2468        })
2469    }
2470
2471    async fn store_blob(
2472        &self,
2473        digest: &CacheDigest,
2474        source: &Path,
2475        connection: &mut ConnectionUploads,
2476    ) -> Result<AgentResponse> {
2477        let path = {
2478            let lock = self.write_lock(digest);
2479            let _guard = lock.lock().await;
2480            if let Some(path) = self.find_verified_blob(digest)? {
2481                path
2482            } else {
2483                let path = self.cas.store_file(digest, source)?;
2484                self.remember_verified_blob(digest, &path);
2485                self.stats.stores.fetch_add(1, Ordering::Relaxed);
2486                self.stats
2487                    .stored_bytes
2488                    .fetch_add(digest.size, Ordering::Relaxed);
2489                path
2490            }
2491        };
2492        // The object is durable locally, so the build has what it needs and the
2493        // remote publication can happen after this request returns.
2494        if let Some(uploads) = &self.uploads {
2495            uploads.queue_blob(digest, path.clone(), connection);
2496        }
2497        Ok(AgentResponse::Stored { path })
2498    }
2499
2500    fn find_verified_blob(&self, digest: &CacheDigest) -> Result<Option<PathBuf>> {
2501        let remembered = self.verified_blobs.lock().unwrap().get(digest).cloned();
2502        if let Some(remembered) = remembered {
2503            // The contents behind this digest were hashed in full when the
2504            // session first reached for them. Hashing again on every later
2505            // lookup would re-read each dependency's rlib once per crate that
2506            // links it, so what is rechecked here is that the file has not
2507            // been written since: an overwrite moves the modification time,
2508            // and truncation or eviction changes the length or removes it.
2509            // Only a replacement that reproduces both -- which no writer in
2510            // the store does, since blobs are published by rename under a
2511            // content-derived name -- would go unnoticed until `mbx cache
2512            // verify` reads the store back in full.
2513            if remembered.is_unchanged() {
2514                return Ok(Some(remembered.path));
2515            }
2516            self.verified_blobs.lock().unwrap().remove(digest);
2517        }
2518        let path = self.cas.find(digest)?;
2519        if let Some(path) = &path {
2520            self.remember_verified_blob(digest, path);
2521        }
2522        Ok(path)
2523    }
2524
2525    fn remember_verified_blob(&self, digest: &CacheDigest, path: &Path) {
2526        let Some(verified) = VerifiedBlob::describe(path) else {
2527            // Without an identity to compare against there is nothing to
2528            // shortcut safely, so the digest stays unremembered and every
2529            // lookup re-reads it.
2530            return;
2531        };
2532        self.verified_blobs
2533            .lock()
2534            .unwrap()
2535            .insert(digest.clone(), verified);
2536    }
2537
2538    async fn find_action_result(&self, action: &CacheDigest) -> Result<AgentResponse> {
2539        if let Some(result) = self.actions.find(action)? {
2540            return Ok(AgentResponse::ActionResult {
2541                result: Some(result),
2542            });
2543        }
2544        if !self.remote_mode.reads() {
2545            return Ok(AgentResponse::ActionResult { result: None });
2546        }
2547        let Some(remote) = &self.remote else {
2548            return Ok(AgentResponse::ActionResult { result: None });
2549        };
2550        let lock = self.action_lock(action);
2551        let _guard = lock.lock().await;
2552        if let Some(result) = self.actions.find(action)? {
2553            return Ok(AgentResponse::ActionResult {
2554                result: Some(result),
2555            });
2556        }
2557        if let Some(result) = self
2558            .pending_remote_actions
2559            .lock()
2560            .unwrap()
2561            .get(action)
2562            .cloned()
2563        {
2564            return Ok(AgentResponse::ActionResult {
2565                result: Some(result),
2566            });
2567        }
2568        let _permit = self.remote_transfers.acquire().await?;
2569        match self.get_remote_action_result(remote, action).await {
2570            Ok(Some(result)) => {
2571                self.pending_remote_actions
2572                    .lock()
2573                    .unwrap()
2574                    .insert(action.clone(), result.clone());
2575                Ok(AgentResponse::ActionResult {
2576                    result: Some(result),
2577                })
2578            }
2579            Ok(None) => Ok(AgentResponse::ActionResult { result: None }),
2580            Err(error) => {
2581                self.note_remote_failure();
2582                warn!(
2583                    "remote cache action lookup failed for {}: {error}",
2584                    action.hash
2585                );
2586                Ok(AgentResponse::ActionResult { result: None })
2587            }
2588        }
2589    }
2590
2591    async fn store_action_result(
2592        &self,
2593        result: &RemoteActionResult,
2594        connection: &ConnectionUploads,
2595    ) -> Result<AgentResponse> {
2596        let path = self.actions.store(result)?;
2597        if let Some(uploads) = &self.uploads {
2598            uploads.queue_action_result(result, connection);
2599        }
2600        Ok(AgentResponse::ActionStored { path })
2601    }
2602
2603    /// Record that a remote operation failed and the build carried on without it.
2604    fn note_remote_failure(&self) {
2605        self.stats.remote_failures.fetch_add(1, Ordering::Relaxed);
2606    }
2607
2608    async fn get_remote_action_result(
2609        &self,
2610        remote: &RemoteCacheClient,
2611        action: &CacheDigest,
2612    ) -> Result<Option<RemoteActionResult>> {
2613        self.stats
2614            .remote_action_lookups
2615            .fetch_add(1, Ordering::Relaxed);
2616        let _timer = AtomicDurationTimer::start(&self.stats.remote_action_lookup_duration_ns);
2617        remote.get_action_result(action).await
2618    }
2619
2620    fn record_action_hit(
2621        &self,
2622        action: &CacheDigest,
2623        restore: RestoreStats,
2624        crate_name: Option<String>,
2625    ) -> Result<AgentResponse> {
2626        validate_crate_name(crate_name.as_deref())?;
2627        if self.actions.find(action)?.is_none() {
2628            let pending = self.pending_remote_actions.lock().unwrap().remove(action);
2629            if let Some(result) = pending {
2630                self.actions.store(&result)?;
2631            } else {
2632                bail!("cannot record a hit for a missing action result");
2633            }
2634        }
2635        self.record_restore(restore);
2636        self.stats.hits.fetch_add(1, Ordering::Relaxed);
2637        self.emit(|| AgentEvent::ActionHit {
2638            crate_name,
2639            restore,
2640        });
2641        Ok(AgentResponse::ActionHitRecorded)
2642    }
2643
2644    fn record_restore(&self, restore: RestoreStats) {
2645        self.record_materialization(restore);
2646        atomic_saturating_add(
2647            &self.stats.avoided_compiler_duration_ns,
2648            restore.avoided_compiler_duration_ns,
2649        );
2650        atomic_saturating_add(&self.stats.restored_output_files, restore.output_files);
2651        atomic_saturating_add(&self.stats.restored_output_bytes, restore.output_bytes);
2652        atomic_saturating_add(
2653            &self.stats.reflinked_output_files,
2654            restore.reflinked_output_files,
2655        );
2656        atomic_saturating_add(
2657            &self.stats.reflinked_output_bytes,
2658            restore.reflinked_output_bytes,
2659        );
2660        atomic_saturating_add(&self.stats.copied_output_files, restore.copied_output_files);
2661        atomic_saturating_add(&self.stats.copied_output_bytes, restore.copied_output_bytes);
2662        atomic_saturating_add(&self.stats.reused_output_files, restore.reused_output_files);
2663        atomic_saturating_add(&self.stats.reused_output_bytes, restore.reused_output_bytes);
2664    }
2665
2666    fn record_compiler_invocation(
2667        &self,
2668        outcome: &str,
2669        crate_name: Option<&str>,
2670        duration_ns: u64,
2671    ) -> Result<AgentResponse> {
2672        if !matches!(
2673            outcome,
2674            "miss" | "unconsulted" | "bypass" | "verification" | "incremental"
2675        ) {
2676            bail!("invalid compiler invocation outcome");
2677        }
2678        validate_crate_name(crate_name)?;
2679        let mut compiler = self.stats.compiler.lock().unwrap();
2680        let stats = compiler.entry(outcome.to_string()).or_default();
2681        stats.invocations = stats.invocations.saturating_add(1);
2682        stats.duration_ns = stats.duration_ns.saturating_add(duration_ns);
2683        drop(compiler);
2684        if outcome != "verification"
2685            && let Some(crate_name) = crate_name.filter(|name| !name.is_empty())
2686        {
2687            let mut slow = self.stats.slow_compilations.lock().unwrap();
2688            let duration = slow.entry(crate_name.to_string()).or_default();
2689            *duration = duration.saturating_add(duration_ns);
2690        }
2691        self.emit(|| AgentEvent::CompilerInvocation {
2692            outcome: outcome.to_string(),
2693            crate_name: crate_name.map(str::to_string),
2694            duration_ns,
2695        });
2696        Ok(AgentResponse::CompilerInvocationRecorded)
2697    }
2698
2699    fn record_materialization(&self, restore: RestoreStats) {
2700        atomic_saturating_add(&self.stats.materialization_duration_ns, restore.duration_ns);
2701    }
2702
2703    /// Surface a shim diagnostic on this process's stderr.
2704    ///
2705    /// The agent lives in the process that owns the session, so printing here
2706    /// reaches the terminal running the build rather than the stderr of the
2707    /// compilation the shim stands in for -- which build scripts read as part
2708    /// of the compiler's answer. Deduplicated because one cause tends to fire
2709    /// once per compilation, and capped so a message unique per compilation
2710    /// cannot scroll the build away.
2711    fn record_warning(&self, message: String) -> Result<AgentResponse> {
2712        if message.is_empty()
2713            || message.len() > MAX_WARNING_BYTES
2714            || message.contains(['\n', '\r', '\0'])
2715        {
2716            bail!("invalid shim warning");
2717        }
2718        let mut warnings = self.warnings.lock().unwrap();
2719        if !warnings.contains(&message) && warnings.len() < MAX_WARNINGS {
2720            eprintln!("mbx[warning]: {message}");
2721            self.emit(|| AgentEvent::Warning {
2722                message: message.clone(),
2723            });
2724            warnings.insert(message);
2725        }
2726        Ok(AgentResponse::WarningRecorded)
2727    }
2728
2729    /// Answer file-digest lookups from the session ledger.
2730    ///
2731    /// A recorded digest is returned only when the requested identity matches
2732    /// the recorded one exactly. The caller statted the file to build the
2733    /// identity, so this is the [`VerifiedBlob`] freshness check with the
2734    /// stat moved to the side that was already making it.
2735    fn find_file_digests(
2736        &self,
2737        scope: FileDigestScope,
2738        files: Vec<FileIdentity>,
2739    ) -> Result<AgentResponse> {
2740        if files.len() > MAX_FILE_DIGEST_BATCH {
2741            bail!("too many file-digest lookups in one request");
2742        }
2743        let ledger = self.file_digests.lock().unwrap();
2744        let digests = files
2745            .into_iter()
2746            .map(|file| {
2747                let recorded = ledger.get(&(scope, file.path.clone()))?;
2748                (recorded.file == file).then(|| recorded.digest.clone())
2749            })
2750            .collect();
2751        Ok(AgentResponse::FileDigests { digests })
2752    }
2753
2754    /// Record digests of files a shim read in full, for later reuse.
2755    ///
2756    /// Capped rather than evicted: a session that outgrows the cap loses the
2757    /// shortcut for the overflow and nothing else, and no workload observed so
2758    /// far comes near it.
2759    fn record_file_digests(
2760        &self,
2761        scope: FileDigestScope,
2762        entries: Vec<RecordedFileDigest>,
2763    ) -> Result<AgentResponse> {
2764        if entries.len() > MAX_FILE_DIGEST_BATCH {
2765            bail!("too many file-digest records in one request");
2766        }
2767        for entry in &entries {
2768            if !entry.file.path.is_absolute() {
2769                bail!("file-digest records need absolute paths");
2770            }
2771            entry.digest.validate()?;
2772            if entry.file.len != entry.digest.size {
2773                bail!("file-digest record length does not match its digest");
2774            }
2775        }
2776        let mut ledger = self.file_digests.lock().unwrap();
2777        for entry in entries {
2778            if ledger.len() >= MAX_FILE_DIGEST_ENTRIES
2779                && !ledger.contains_key(&(scope, entry.file.path.clone()))
2780            {
2781                break;
2782            }
2783            ledger.insert((scope, entry.file.path.clone()), entry);
2784        }
2785        Ok(AgentResponse::FileDigestsRecorded)
2786    }
2787
2788    fn find_action_prediction(
2789        &self,
2790        task: &str,
2791        invocation: &CacheDigest,
2792    ) -> Result<AgentResponse> {
2793        validate_task_identity(task)?;
2794        invocation.validate()?;
2795        let prediction = self
2796            .task_actions
2797            .lock()
2798            .unwrap()
2799            .get(task)
2800            .and_then(|state| state.predictions.get(invocation))
2801            .cloned();
2802        Ok(AgentResponse::ActionPrediction { prediction })
2803    }
2804
2805    fn record_action_prediction(
2806        &self,
2807        task: &str,
2808        prediction: ActionPrediction,
2809    ) -> Result<AgentResponse> {
2810        validate_task_identity(task)?;
2811        prediction.validate()?;
2812        let mut tasks = self.task_actions.lock().unwrap();
2813        let state = tasks.entry(task.to_string()).or_default();
2814        if !state.predictions.contains_key(&prediction.invocation)
2815            && state.predictions.len() >= MAX_TASK_ACTION_PREDICTIONS
2816        {
2817            bail!("task action manifest contains too many predictions");
2818        }
2819        state
2820            .predictions
2821            .insert(prediction.invocation.clone(), prediction.clone());
2822        state
2823            .pending_predictions
2824            .insert(prediction.invocation.clone(), prediction);
2825        Ok(AgentResponse::ActionPredictionRecorded)
2826    }
2827
2828    fn executable_identity_key(
2829        &self,
2830        executable: PathBuf,
2831        environment: BTreeMap<String, Option<String>>,
2832    ) -> Result<ExecutableIdentityKey> {
2833        // Restricted to the variables that actually select what an identity
2834        // probe reports: the toolchain rustup resolves, and the SDK a linker
2835        // driver builds against. Anything else would let one key stand for two
2836        // different compilers.
2837        if !environment.keys().all(|name| {
2838            matches!(
2839                name.as_str(),
2840                "RUSTUP_HOME" | "RUSTUP_TOOLCHAIN" | "SDKROOT" | "MACOSX_DEPLOYMENT_TARGET"
2841            )
2842        }) {
2843            bail!("executable identity contains an unsupported environment variable");
2844        }
2845        Ok(ExecutableIdentityKey {
2846            executable,
2847            environment,
2848        })
2849    }
2850
2851    fn find_executable_identity(
2852        &self,
2853        executable: PathBuf,
2854        environment: BTreeMap<String, Option<String>>,
2855    ) -> Result<AgentResponse> {
2856        let key = self.executable_identity_key(executable, environment)?;
2857        let stdout = self
2858            .executable_identities
2859            .lock()
2860            .unwrap()
2861            .get(&key)
2862            .cloned();
2863        Ok(AgentResponse::ExecutableIdentity { stdout })
2864    }
2865
2866    fn store_executable_identity(
2867        &self,
2868        executable: PathBuf,
2869        environment: BTreeMap<String, Option<String>>,
2870        stdout: Vec<u8>,
2871    ) -> Result<AgentResponse> {
2872        if stdout.len() > MAX_EXECUTABLE_IDENTITY_SIZE {
2873            bail!("executable identity exceeds {MAX_EXECUTABLE_IDENTITY_SIZE} bytes");
2874        }
2875        let key = self.executable_identity_key(executable, environment)?;
2876        let mut identities = self.executable_identities.lock().unwrap();
2877        let is_new = !identities.contains_key(&key);
2878        let previous_size = identities.get(&key).map_or(0, Vec::len);
2879        if is_new && identities.len() >= MAX_EXECUTABLE_IDENTITIES {
2880            bail!("executable identity cache contains too many entries");
2881        }
2882        let retained_bytes = identities.values().map(Vec::len).sum::<usize>();
2883        if retained_bytes - previous_size + stdout.len() > MAX_EXECUTABLE_IDENTITY_BYTES {
2884            bail!("executable identity cache contains too many bytes");
2885        }
2886        identities.insert(key, stdout.clone());
2887        Ok(AgentResponse::ExecutableIdentity {
2888            stdout: Some(stdout),
2889        })
2890    }
2891
2892    /// Serve newline-delimited protocol requests on an authenticated session stream.
2893    pub async fn handle_connection<S>(&self, stream: S) -> Result<()>
2894    where
2895        S: AsyncRead + AsyncWrite + Unpin,
2896    {
2897        let (reader, mut writer) = tokio::io::split(stream);
2898        let mut reader = BufReader::new(reader);
2899        let hello = read_request(&mut reader)
2900            .await?
2901            .ok_or_else(|| eyre::eyre!("connection closed before the agent handshake"))?;
2902        let request: AgentRequest = serde_json::from_str(&hello)?;
2903        match request {
2904            AgentRequest::Hello {
2905                protocol,
2906                client_version,
2907            } if protocol == AGENT_PROTOCOL_VERSION && client_version == self.version.as_ref() => {}
2908            AgentRequest::Hello { protocol, .. } if protocol != AGENT_PROTOCOL_VERSION => {
2909                send_response(
2910                    &mut writer,
2911                    &AgentResponse::Error {
2912                        message: format!(
2913                            "unsupported agent protocol {protocol}; expected {AGENT_PROTOCOL_VERSION}"
2914                        ),
2915                    },
2916                )
2917                .await?;
2918                return Ok(());
2919            }
2920            AgentRequest::Hello { client_version, .. } => {
2921                send_response(
2922                    &mut writer,
2923                    &AgentResponse::Error {
2924                        message: format!(
2925                            "cache client {client_version} does not match agent {}",
2926                            self.version
2927                        ),
2928                    },
2929                )
2930                .await?;
2931                return Ok(());
2932            }
2933            _ => bail!("the first agent request must be hello"),
2934        }
2935        send_response(
2936            &mut writer,
2937            &AgentResponse::Hello {
2938                protocol: AGENT_PROTOCOL_VERSION,
2939                agent_version: self.version.to_string(),
2940            },
2941        )
2942        .await?;
2943
2944        // Tickets accumulate for the life of the connection because a shim
2945        // publishes a compilation's blobs and the action result naming them over
2946        // one connection, in that order.
2947        let mut connection = ConnectionUploads::default();
2948        while let Some(line) = read_request(&mut reader).await? {
2949            let response = match serde_json::from_str(&line) {
2950                Ok(request) => self.respond_on(request, &mut connection).await,
2951                Err(error) => AgentResponse::Error {
2952                    message: format!("invalid agent request: {error}"),
2953                },
2954            };
2955            send_response(&mut writer, &response).await?;
2956        }
2957        Ok(())
2958    }
2959}
2960
2961/// Read one newline-delimited request, refusing one that grows past the cap.
2962///
2963/// Any process running as this user can open the session socket, so a request
2964/// that never terminates its line must not be able to grow the agent's memory
2965/// without bound.
2966async fn read_request<R>(reader: &mut R) -> Result<Option<String>>
2967where
2968    R: AsyncBufRead + Unpin,
2969{
2970    let mut line = Vec::new();
2971    loop {
2972        let available = reader.fill_buf().await?;
2973        if available.is_empty() {
2974            break;
2975        }
2976        let (consumed, complete) = match available.iter().position(|byte| *byte == b'\n') {
2977            Some(index) => (index, true),
2978            None => (available.len(), false),
2979        };
2980        if line.len() + consumed > MAX_REQUEST_BYTES {
2981            bail!("agent request exceeded {MAX_REQUEST_BYTES} bytes");
2982        }
2983        line.extend_from_slice(&available[..consumed]);
2984        // The newline itself is consumed but never kept.
2985        reader.consume(consumed + usize::from(complete));
2986        if complete {
2987            return Ok(Some(String::from_utf8(line)?));
2988        }
2989    }
2990    if line.is_empty() {
2991        Ok(None)
2992    } else {
2993        Ok(Some(String::from_utf8(line)?))
2994    }
2995}
2996
2997/// Whether `task` is a well-formed task action identity.
2998///
2999/// Identities name files and directories in the store, so anything that reads
3000/// the store back has to be able to tell an identity from whatever else a user
3001/// left lying there.
3002pub fn is_task_identity(task: &str) -> bool {
3003    task.len() == 64
3004        && task
3005            .bytes()
3006            .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
3007}
3008
3009fn validate_task_identity(task: &str) -> Result<()> {
3010    if !is_task_identity(task) {
3011        bail!("invalid task action identity");
3012    }
3013    Ok(())
3014}
3015
3016/// Where a store keeps its task prediction manifests.
3017fn task_manifest_dir(store: &Path) -> PathBuf {
3018    store.join("task-manifests").join("v1")
3019}
3020
3021/// The action digests a task's prediction manifest recorded.
3022///
3023/// Read straight off disk rather than through an agent, because a collector
3024/// needs the action set of tasks no session is running. A manifest that is
3025/// missing or no longer parseable yields no actions rather than an error: this
3026/// is a prediction index, so the worst a thin answer costs is a cold prefetch,
3027/// or an object collected earlier than it deserved.
3028pub fn task_manifest_actions(store: &Path, task: &str) -> Result<Vec<CacheDigest>> {
3029    validate_task_identity(task)?;
3030    let path = task_manifest_dir(store).join(format!("{task}.json"));
3031    let bytes = match fs::read(&path) {
3032        Ok(bytes) => bytes,
3033        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
3034        Err(error) => {
3035            return Err(error).wrap_err_with(|| format!("failed to read {}", path.display()));
3036        }
3037    };
3038    let Ok(manifest) = serde_json::from_slice::<TaskActionManifest>(&bytes) else {
3039        return Ok(Vec::new());
3040    };
3041    if validate_task_manifest(&manifest, task).is_err() {
3042        return Ok(Vec::new());
3043    }
3044    Ok(manifest
3045        .predictions
3046        .into_iter()
3047        .map(|prediction| prediction.action)
3048        .collect())
3049}
3050
3051fn validate_task_manifest(manifest: &TaskActionManifest, task: &str) -> Result<()> {
3052    if manifest.task == task && manifest.validate() {
3053        Ok(())
3054    } else {
3055        bail!("invalid task action manifest")
3056    }
3057}
3058
3059fn merge_task_manifests(
3060    task: &str,
3061    base: Option<TaskActionManifest>,
3062    update: TaskActionManifest,
3063) -> Result<TaskActionManifest> {
3064    validate_task_manifest(&update, task)?;
3065    let mut predictions = BTreeMap::new();
3066    if let Some(base) = base {
3067        validate_task_manifest(&base, task)?;
3068        predictions.extend(
3069            base.predictions
3070                .into_iter()
3071                .map(|prediction| (prediction.invocation.clone(), prediction)),
3072        );
3073    }
3074    predictions.extend(
3075        update
3076            .predictions
3077            .into_iter()
3078            .map(|prediction| (prediction.invocation.clone(), prediction)),
3079    );
3080    let manifest = TaskActionManifest {
3081        version: TASK_ACTION_MANIFEST_VERSION,
3082        task: task.to_owned(),
3083        predictions: predictions.into_values().collect(),
3084    };
3085    validate_task_manifest(&manifest, task)?;
3086    Ok(manifest)
3087}
3088
3089fn merge_remote_task_manifest(
3090    task: &str,
3091    remote: TaskActionManifest,
3092    local: TaskActionManifest,
3093) -> (TaskActionManifest, bool) {
3094    match merge_task_manifests(task, Some(remote), local.clone()) {
3095        Ok(manifest) => (manifest, true),
3096        Err(error) => {
3097            warn!("remote task action manifest merge failed for {task}: {error}");
3098            (local, false)
3099        }
3100    }
3101}
3102
3103async fn send_response(
3104    writer: &mut (impl AsyncWrite + Unpin),
3105    response: &AgentResponse,
3106) -> Result<()> {
3107    let mut encoded = serde_json::to_vec(response)?;
3108    encoded.push(b'\n');
3109    writer.write_all(&encoded).await?;
3110    writer.flush().await?;
3111    Ok(())
3112}
3113
3114#[cfg(test)]
3115#[path = "agent_tests.rs"]
3116mod tests;